Beyond COVID-19 Diagnosis: Prognosis with Hierarchical Graph Representation Learning

1 INTRODUCTION . Coronavirus disease 2019 ( COVID-19 ) has resulted in an ongoing pandemic in the world . To control the sources of infection and cut off the channels of transmission , rapid testing and detection are of vital importance . The reverse transcription polymerase chain reaction ( RT-PCR ) is a widely-used screening technology and viewed as the standard method for suspected cases . However , this method highly relies upon the required lab facilities and the diagnostic kits . In addition , the sensitivity of RT-PCR is not high enough for early diagnosis ( Ai et al. , 2020 ; Fang et al. , 2020 ) . To mitigate the limitations of RT-PCR , the computed tomography ( CT ) has been widely used as an effective complementary method , which can provide medical images of the lung area to reveal the details of the disease and its prognosis ( Huang et al. , 2020 ; Chung et al. , 2020 ) , for which RT-PCR can not . Additionally , CT has also been proved to be useful in monitoring the COVID-19 disease progression and the therapeutic efficacy evaluation ( Rodriguez-Morales et al. , 2020 ; Liechti et al. , 2020 ) . The chest CT slices of a patient have a sequential and hierarchical data structure . The relationship between slices possesses more information than the order of the slices . The adjacent ones with the same abnormality could be considered as one lesion . The slices containing the same type of lesions may not be continuous as the lesions are distributed in various lung parts . We propose a diagnosis and prognosis system that combines graph convolutional networks ( GCNs ) and a distance aware pooling , which integrates the information from all slices in the chest CT scans for optimal decision making . Our major contributions are three-fold : ( 1 ) Owing to the sequential structure of CT images , this is the first work to utilize GCNs to extract node information hierarchically , and conduct both diagnosis and prognosis for COVID-19 . The prognosis can help facilitate medical resources , e.g. , ventilators or admission to Intensive Care Units ( ICUs ) , more efficiently by triaging mild or severe patients . ( 2 ) A novel pooling method called distance aware pooling , is proposed to aggregate the graph , i.e. , the patient ’ s CT scan , effectively . The new pooling method integrated with GCNs can aggregate a densely connected graph efficiently . ( 3 ) The new model can localize the most informative slices within a chest CT scan , which significantly reduces the amount of work for radiologists . 2 RELATED WORK . AI-assisted and CT-based COVID-19 Diagnosis and Prognosis . Although RT-PCR is the standard way for COVID-19 diagnosis , there are many limitations using RT-PCR along , e.g. , time delay in receiving an RT-PCR test , occurrence of false negatives , and no prognostic information provided , etc . CT images are often recommended as an alternative for precise lesion detection ( Alizadehsani et al. , 2020 ) . However , as each CT scan includes a large number of ( up to several hundreds of ) image slices , it requires much time and labor of the radiologists ( Shoeibi et al. , 2020 ) . Furthermore , since the radiological appearances of COVID-19 are similar to other types of pneumonia , radiologists need to go through extensive training before they can achieve high diagnostic accuracy ( Shi et al. , 2020 ) . Recently , several AI-assisted and CT-based COVID-19 diagnostic systems have been developed . Chen et al . ( 2020 ) use Unet++ ( Zhou et al. , 2018 ) to segment infectious areas in the lung . Butt et al . ( 2020 ) develop a deep learning model to detect lesions from CT images , and then use 3D ResNet to classify the images into COVID-19 , influenza-A viral pneumonia or healthy groups . Song et al . ( 2020 ) use the whole lung for diagnosis instead of only extracting the lesions . Wang et al . ( 2020 ) propose an AI system that can diagnose COVID-19 patients as well as conducting the prognostic analysis . Graph Neural Networks . Various types of graph neural networks ( GNNs ) have been proposed , which can be divided into spectral or non-spectral domains . In the spectral domain , the Fourier transformation and graph Laplacian define the convolutional filters ( Bruna et al. , 2013 ) . By utilizing Chebyshev polynomials , Kipf & Welling ( 2016 ) simplify the filters of graph convolution , rendering a layer-wise propagation method . However , the generalization of spectral methods may not be ideal due to the variety of graphs ( Bronstein et al. , 2017 ) . Non-spectral methods focus on the local topology of nodes , directly working on graphs instead of the Fourier domain . Methods proposed by Hamilton et al . ( 2017 ) , Monti et al . ( 2017 ) and Veličković et al . ( 2017 ) aggregate nodes based on adjacent nodes when the next layer is created . This aggregation process , as mentioned by Gilmer et al . ( 2017 ) , can be regarded as a message-passing process . Pooling Methods Pooling methods allow GNNs to hierarchically aggregate nodes , obtaining and assembling local information of graphs . The major purpose of the hierarchical pooling method is to use a locally based model to aggregate nodes in each layer , so that a higher level graph representation can be created ( Lee et al. , 2019 ) . The self-attention based pooling method ( Mao et al. , 2018 ) is implemented for video classification , locally obtaining weighted and fused feature sequences . Spectral pooling methods , such as the one proposed by Ma et al . ( 2019 ) , focus on the application of eigen-decomposition to capture the graph information . However , since spectral pooling methods are computationally demanding , they may not be accessible to large graphs . Non-spectral pooling methods are scalable to large graphs and pay more attention to the local structures of graphs . Adaptive Structure Aware Pooling ( ASAP ) ( Ranjan et al. , 2020 ) and DiffPool ( Ying et al. , 2018 ) resort to clustering techniques , aggregating nodes into different clusters , and then choosing top clusters based on cluster ranking scores ( Gao & Ji , 2019 ) . However , the ASAP method partially ignores the edge weight information when using hops to aggregate nodes , resulting in unstable convergence . 3 METHODOLOGY . We propose a GCN-based diagnosis and prognosis method that models the sequential slices of CT scans hierarchically . To downsample and learn graph-level representation from the input node features , a novel distance aware pooling method is proposed . In this paper , the node features refer to the slices in a CT scan . The model gradually extracts information from the slice level to the patient level by graph convolution and pooling . Eventually , a higher-level representation is learned , and further used for diagnosis , prognosis , and lesion localization . The schema of our model is illustrated in Figure 1 , which is composed of GCNs , pooling modules , a multilayer perceptron ( MLP ) classifier , and a one-drop localization module . The graph convolution-based method can integrate all slices in the chest CT scans for optimal decision making . Furthermore , we propose the one-drop localization to localize the most informative slices , so that radiologists may focus on those recommended slices with the most suspected lesion areas . Consequently , the proposed model can produce visual explanations for the diagnosis and prognosis , making the decision more transparent and explainable . We argue that this method could effectively assist radiologists by reducing redundancies in the vast amount of CT slices during diagnosis and prognosis in the clinical settings . 3.1 PROBLEM STATEMENT . Let G ( V , E ) be a patient ’ s CT scan graph , with |V| = N nodes and |E| edges , where |·| represents the cardinality of a set . For each vi ∈ V , xi is the corresponding d-dimensional vector . Let X ∈ RN×d be the node feature matrix , and Aadj ∈ RN×N be the weighted adjacency matrix . Each entry in Aadj is defined based on cosine similarity , which is Aadji , j = < xi , xj > / ( ‖xi‖ · ‖xj‖ ) . Define the distance matrix Adis as Adis = 1 − Aadj , where 1 is the matrix with all entries of 1 ’ s . The definition of Adis means that the more similar two nodes are , the shorter distance they have . In the construction of Aadj , the diagonal entries are automatically 1 ’ s since i = j is allowed . Each graph G has a label y . For diagnosis , the label represents its class from normal , common pneumonia , and COVID-19 . For prognosis , the class indicates whether a COVID-19 positive patient develops into sever/critical illness status . Thus , the diagnosis and prognosis of COVID-19 is a task of graph classification under our setting . Given a training dataset T = { ( G1 , y1 ) , . . . , ( GM , yM ) } , the goal is to learn a mapping f : G → y , which classifies a graph G into the corresponding class y . Our model is composed of two modules : f1 includes node convolution and feature pooling , and f2 involves a MLP classifier which determines the class of each graph . At each layer of f1 , node embeddings and cluster membership are learnt iteratively . The first module can be written as f1 : G → Gp , where Gp is the pooled graph with fewer nodes and a hierarchical feature representation . The second module is f2 : Gp → y , which utilizes the graph-level representation learnt for patient diagnosis and prognosis . The two modules are integrated in an end-to-end fashion . That is to say , the cluster assignment is learnt merely based on the graph classification objective . 3.2 NODE CONVOLUTION AND FEATURE POOLING . 3.2.1 NODE CONVOLUTION . Node convolution applies the graph convolutional network to obtain a high-level node feature representation in the feature matrix X . Although several methods exist to construct the convolutional network , the method recommended by Kipf & Welling ( 2016 ) is effective for our case , which is given by X ( l+1 ) = σ ( √ D ( l ) Aadj ( l ) √ D ( l ) X ( l ) W ( l ) ) , where D ( l ) is the diagonal degree matrix of Aadj ( l ) − I , and W ( l ) ∈ Rd×k is learnable weight matrix at the l-th layer . Due to the application of feature pooling , the topology of the graph changes at each layer , and thus the dimension of matrices involved are reduced accordingly . 3.2.2 DISTANCE AWARE POOLING METHOD . We propose an innovative pooling method , which includes graph-based clustering and feature pooling . Below , we outline the pooling module and illustrate how it is integrated into an end-to-end GCN based model . Empirically , it is shown to be more robust for densely connected graphs . The overall structure of the pooling method is illustrated in figure 2 . Improved receptive field . The concept of the receptive field , RF , used in the convolutional neural network ( CNN ) was extended to GNNs ( Ranjan et al. , 2020 ) . They defined RFnode as the number of hops required to cover the neighborhood of a given node , such that given a chosen node , a cluster can be obtained based on a fixed receptive field h. However , this design may not be applicable to densely connected graphs , because one node may be connected to most of the nodes in the graph . Even given a small value of h , for instance , h = 1 which is the default value ( Ranjan et al. , 2020 ) , clusters formed within hop h = 1 may include most of the nodes of the graph . Hence , the clustering step may be inefficient . In addition , the value of hops can only be integers , which loses the edge weight information in the clustering process . Therefore , we define an improved receptive field for densely connected graphs RF d denoted by hd , which can be treated as a radius centered at a given node . The value of hd is not restricted to integers but can be a positive real number . Define N ( vi ) as the local neighborhood of the node vi , and Nhd ( vi ) as the RF d neighborhood of the node vi with a radius hd , and ∀vj ∈ N ( vi ) , vi and vj are connected by the edge ( vi , vj ) . Node clustering . Inspired by the clustering and ranking ideas mentioned in Ranjan et al . ( 2020 ) and Gao & Ji ( 2019 ) , we propose a local node clustering , and score ranking method . Each node is considered as a center of a cluster for a given hd . Then , we score all the clusters and choose the top k proportion of them to represent the next layer ’ s nodes with pooled feature values , where k is a hyperparameter . Clustering ranking . Given a node vi and a radius hd , Nhd ( vi ) is the corresponding RF d neighborhood . Let I1 ( vi ) be the index set of the nodes in Nhd ( vi ) . The cluster score is defined as αi = ∑ m , n∈I1 ( vi ) A dis m , n/|Nhd ( vi ) | , where m 6= n. If αi is small , nodes in Nhd ( vi ) are close to each other . The top k proportion of clusters form the next layer ’ s nodes . Selecting cluster centers . We first introduce some notation : ∀vj ∈ Nhd ( vi ) , define Vj = N [ vj ] ∩ Nhd ( vi ) as the set of nodes connected to vj in Nhd ( vi ) , where N [ vj ] is the closed neighborhood of the node vj . Let I2 ( vj ) be the index set of the nodes in Vj . The node score is defined as bj =∑ c∈ I2 ( vi ) A dis j , c /|Vj | . The node with the smallest value of bj is chosen as the center of the cluster , and is used to represent a new node in the next layer ’ s node representation . Center node feature pooling . Based on the node scores , we rank the nodes in Vj and assign a weight wr to xr , where r ∈ I2 ( vj ) . Let b̂ = 1 − b . The weight vector w is defined as wr = b̂r/ ( ∑ i∈I2 ( vj ) b̂i + ) , where 1 is a vector with all 1 ’ s , is an extremely small positive value to avoid 0 ’ s in the denominator , and b is the vector containing all the weights of nodes in Vj . The value of the pooled center feature is defined as xp = Xrw , and then xp is used as the center feature representing Nhd ( vi ) . Next layer node connectivity . Following the idea in Ying et al . ( 2018 ) , the connectivity of the nodes in the next layer is preserved as follows . According to the above ranking and pooling methods , a pooled graph Gp with the node set Vp is obtained . The next step is to decide the pooled adjacency matrix Aadj , p . Define matrix S such that the columns of S are the top k clusters ’ weight vectors w. Hence , the pooled adjacency matrix is defined as Aadj , p = STAadjS .

Summary: The manuscript proposes  a distance aware pooling method to use in  graph convolutional neural for predicting whether a subject is infected with Covid-19 (diagnosis) and progression of the disease (prognosis). Experiments were conducted on CT images from three groups: Covid-19 group, common pneumonia group, and heathy group with about 900 samples in each group. The proposed model achieved 94.7% accuracy.