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 paper is an application of GCN with good features on chest CT scan images for Covid-19 diagnosis and prognosis. First of all this is a relevant and appreciated effort when the world is fighting the pandemic. Hence some bonus points is directed towards that. As a whole, to the representation learning community, it adds limited research values apart from being an application of GCN which is aligned to the application track of ICLR. The paper claims that with less than 1% number of total parameters in the baseline 3D ResNet model, their method achieves 94.7% accuracy for diagnosis, which is marginally better than state of art - however whether the model was over-fitted is not clear. Prognosis information is an added claim, though automation part is not integrated.