B EYOND COVID-19 D IAGNOSIS : P ROGNOSIS WITH H IERARCHICAL G RAPH R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Coronavirus disease 2019 (COVID-19), the pandemic that is spreading fast globally, has caused over 34 million conﬁrmed cases.
Apart from the reverse transcription polymerase chain reaction (RT-PCR) the chest computed tomography (CT) is viewed as a standard and effective tool for disease diagnosis and progression monitoring.
We propose a diagnosis and prognosis model based on graph con volutional networks
GCNs).
The chest CT scan of a patient typically involving hundreds of sectional images in sequential order is formulated as a densely con- nected weighted graph.
A novel distance aware pooling is proposed to abstract the node information hierarchically which is robust and efﬁcient for such densely connected graphs
Our method combining GCNs and distance aware pooling, can integrate the information from all slices in the chest CT scans for optimal decision making which leads to the state of the art accuracy in the COVID
19 diagnosis and prognosis.
With less than 1% number of total parameters in the baseline
3D
ResNet model our method achieves
8% accuracy for diagnosis
It has a improvement compared with the baseline model on the same dataset.
In addition, we can localize the most informative slices with disease lesions for COVID- within a large sequence of chest CT images.
The proposed model can produce visual explanations for the diagnosis and prognosis making the decision more transparent and explainable while RT-PCR only leads to the test result with no prognosis information.
The prognosis analysis can help hospitals or clinical centers designate medical resources more efﬁciently and better support clinicians to determine the proper clinical treatment.
I NTRODUCTION
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 cannot.
Additionally, CT has also been proved to be useful in monitoring the COVID-19 disease progression and the therapeutic efﬁcacy 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 ﬁrst 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 efﬁciently 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 efﬁciently.
The new model can localize the most informative slices within a chest CT scan, which signiﬁcantly reduces the amount of work for radiologists.
R ELATED W ORK 2. AI- ASSISTED AND CT- BASED COVID-19 D IAGNOSIS AND P ROGNOSIS
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, inﬂuenza-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.
2. G RAPH N EURAL N ETWORKS
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 deﬁne the convolutional ﬁlters (Bruna et al., 2013).
By utilizing Chebyshev polynomials,
Kipf & Welling (2016) simplify the ﬁlters 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.
Monti et al.
and Veli ˇ ckovi ´ c 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.
2. P OOLING M ETHODS
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 classiﬁcation, 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.
M ETHODOLOGY
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) classiﬁer, 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. P ROBLEM 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 v i ∈ V , x i is the corresponding d -dimensional vector.
Let X ∈ R N ×
d be the node feature matrix, and A adj ∈ R N × N be the weighted adjacency matrix.
Each entry in
A adj is deﬁned based on cosine similarity, which is A adji,j = < x i , x j >/ ( (cid:107) x
i (cid:107) · (cid:107) x j (cid:107) ) .
Deﬁne the distance matrix A dis as A dis = 1 − A adj , where 1 is the matrix with all entries of 1’s.
The deﬁnition of A dis means that the more similar two nodes are, the shorter distance they have.
In the construction of A adj , 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 classiﬁcation under our setting.
Given a training dataset T = { ( G 1 , y 1 ) , . . .
, ( G M , y M ) } , the goal is to learn a mapping f : G → y , which classiﬁes a graph G into the corresponding class y .
Our model is composed of two modules: f 1 includes node convolution and feature pooling, and f 2 involves a MLP classiﬁer which determines the class of each graph.
At each layer of f 1 , node embeddings and cluster membership are learnt iteratively.
The ﬁrst module can be written as f 1 :
G → G p , where G p is the pooled graph with fewer nodes and a hierarchical feature representation.
The second module is f 2 : G p → 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 classiﬁcation objective.
3. N ODE CONVOLUTION AND FEATURE POOLING 3.2. N ODE 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)
=
σ ( (cid:112) D ( l ) A adj ( l ) (cid:112) D ( l ) X
( l ) W ( l ) ) , where D ( l ) is the diagonal degree matrix of A adj ( l )
− I , and W ( l ) ∈ R
d × 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. D ISTANCE A WARE P OOLING M ETHOD
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 ﬁgure 2.
Improved receptive ﬁeld.
The concept of the receptive ﬁeld, RF , used in the convolutional neural network (CNN) was extended to GNNs (Ranjan et al., 2020).
They deﬁned RF node 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 ﬁxed receptive ﬁeld 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 inefﬁcient.
In addition, the value of hops can only be integers, which loses the edge weight information in the clustering process.
Therefore, we deﬁne an improved receptive ﬁeld for densely connected graphs RF d denoted by h d , which can be treated as a radius centered at a given node.
The value of h d is not restricted to integers but can be a positive real number.
Deﬁne N ( v i ) as the local neighborhood of the node v i , and N h d ( v i ) as the RF d neighborhood of the node v i with a radius h d , and ∀ v j ∈ N ( v i ) , v i and v j are connected by the edge ( v i , v j ) .
Node clustering.
Inspired by the clustering and ranking ideas mentioned in Ranjan et al.
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 h d .
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 v i and a radius h d , N h d ( v i ) is the corresponding RF d neighborhood.
Let I 1 ( v i ) be the index set of the nodes in N h d ( v i ) .
The cluster score is deﬁned as α i =
(cid:80) m,n ∈
I 1 ( v i )
A dism,n / |N h
d ( v i ) | , where m (cid:54) = n .
If α i is small, nodes in N h d ( v i ) are close to each other.
The top k proportion of clusters form the next layer’s nodes.
Selecting cluster centers.
We ﬁrst introduce some notation: ∀ v j ∈ N h d ( v i ) , deﬁne V j = N [ v j ] ∩
N h d ( v i ) as the set of nodes connected to v j in N h d ( v i ) , where N [ v j ] is the closed neighborhood of the node v j .
Let I 2 ( v j ) be the index set of the nodes in V j .
The node score is deﬁned as b j = (cid:80) c ∈
I 2 ( v i )
A disj,c / | V j | .
The node with the smallest value of b j 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 V j and assign a weight w r to x r , where r ∈
I 2 ( v j ) .
Let ˆ b = 1 − b .
The weight vector w is deﬁned as w r = ˆ b r / ( (cid:80) i ∈
I 2 ( v j ) ˆ
b i + (cid:15) ) , where 1 is a vector with all 1’s, (cid:15) 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 V j .
The value of the pooled center feature is deﬁned as x p = X r w , and then x p is used as the center feature representing N h d ( v i ) .
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 G p with the node set V p is obtained.
The next step is to decide the pooled adjacency matrix A adj,p .
Deﬁne matrix S such that the columns of S are the top k clusters’ weight vectors w .
Hence, the pooled adjacency matrix is deﬁned as A adj,p = S T A adj S .
3. C LASSIFICATION AND L OCALIZATION B ASED ON P OOLED G RAPHS .
Graph classiﬁcation.
A hierarchical representation for each patient is abstracted by the above GCNs with distance aware pooling.
The representation is an N (cid:48)
× D (cid:48) feature matrix, where N (cid:48) is the number of clusters and D (cid:48) is the number of features in each cluster.
For each feature, the mean and the maximum over the N (cid:48) clusters are calculated.
Subsequently, we obtain a 2 D (cid:48) feature vector.
The ﬁrst D (cid:48) elements are the mean values of each feature, and the second D (cid:48) elements are the maximal values.
Then, we use an MLP to classify the 2 D (cid:48) feature vector into one of the three classes: COVID-19, common pneumonia, or normal.
The MLP classiﬁer is mainly composed of two linear layers.
For each input 2 D (cid:48) feature vector, it can output the scores of each class, normalized by the softmax function.
Weakly supervised informative CT slices localization.
Besides the classiﬁcation, localizing the
CT slices with lesions is also vital for the diagnosis.
To achieve this goal, we propose the onedrop localization to select the most informative CT slices for the model to make a decision.
This method does not require mask annotation for lesion, and thus the localization is learned in a weakly supervised manner.
We are inspired by the stepwise regression that is used for knowledge discovery (Cios et al., 1998), an automatic procedure that integrates the variable selection.
One of the ideas is backward elimination.
For each patient with the N (cid:48) × D (cid:48) feature matrix, we ﬁrst predict the target class using the above MLP classiﬁer.
Then, we drop one cluster each time so that we obtain N (cid:48) new feature matrices in total with an equal size of ( N (cid:48) − 1) × D (cid:48) .
For each of the new feature matrices, the score of the target class is calculated using the same MLP classiﬁer.
The lowest score from the
N (cid:48) results and the corresponding dropped cluster are chosen.
Since ignoring this cluster leads to the lowest score of the target class, this cluster should contain the most crucial diagnosis information.
We can then determine the center of this cluster, and trace back to the corresponding input CT slice.
We further localize the top k s CT slices with the highest similarity to the cluster center.
We believe this whole k s + 1 CT slices are most informative for the diagnosis, which can be suggested as a reference to the radiologists.
E XPERIMENTS AND R ESULTS
For diagnosis and prognosis, we compare our model with the baseline, a 3D ResNet-18 classiﬁcation network (Zhang et al., 2020), and the state-of-the-art graph classiﬁcation method, GCNs with ASAP (Ranjan et al., 2020).
Moreover, we appraise whether the proposed method can deliver meaningful and interpretable clusters on the input chest CT scans by comparing the localization results with the slices containing lesions.
4. D ATASET
We utilize the CT dataset from the 2019 Novel Coronavirus Resource (2019nCoVR) (Zhang et al., 2020).
The dataset includes the complete chest CT scans of 929 COVID-19 positive patients,common pneumonia patients, and 849 healthy individuals.
The dataset also provides these patients’ clinical prognosis, whether the patients developed into severe/critical illness status, referring to the admission to ICU, mechanical ventilation, or death.
The prognosis analysis could support the hospitals to designate medical resources more efﬁciently.
In addition, the dataset summarizes the slices with lesions for COVID-19 positive and common pneumonia patients, which can be used to evaluate the one-drop localization method.
Each CT slice is normalized into the dimension 256 × 256 .
For each chest CT scan, we use systematic sampling to ensure that an equal number of slices are chosen.
4. D ATA P REPROCESSING : C HEST CT S CAN I MAGE F EATURE EXTRACTION
CNN feature extraction.
Feature extraction is a common strategy in computer vision.
In our case, we utilize Inception V3 (Szegedy et al., 2016) pretrained on ImageNet (Russakovsky et al., 2015).
The feature map of the bottleneck layer, which is the last layer before the ﬂatten operation, is regarded as the node representation in the graph.
In our case, each node represents a slice in a chest
CT scan, and the node representation is a vector of dimension 2048 .
Wavelet decomposition extraction.
Considering each slice in CT scans as a 2-dimensional signal, it can be viewed as a function with two variables, which can be reconstructed as a summation of wavelet functions multiplying their coefﬁcients for a given resolution (Mallat, 1989).
The formula required to reconstruct the image is
[Equation]
(cid:88)
[Equation]
j [Equation] [Equation] where f ( x, y ) is the energy function of an image, M and N are the image dimension, W ϕ and W Ψ are the coefﬁcients of scaling function ϕ and wavelet function Ψ respectively.
The coefﬁcient W ϕ is used in down streaming analysis since it contains the major approximation information of an image.
We choose the Haar wavelet function with resolution 3.
The ﬂattened approximation matrix of the image signal, which is a vector of dimension 1024 , is used as the feature embedding of a slice in a
CT scan.
4. I MPLEMENTATION D ETAILS
For dataset 19nCoVR, we use systematic sampling to ensure that 48 number of slices for each CT scan are chosen.
CT scans of 60% individual are randomly chosen as the training set, 25% as the test set, and the remaining 15% for the validation.
To avoid information leakage, the dataset is split according to individuals instead of the CT scan.
Similar to Zhang et al.
and Ying et al.
(2018), we repeat the data splitting aforementioned on 20 random seeds.
For each random seed, the model is trained from scratch.
The maximum, average and standard deviation of test accuracies are reported.
We use ﬁve GCN layers, and the MLP classiﬁer is composed of two fully connected layers.
The pooling proportion k is set as 0.5 for the experiments of diagnosis and prognosis.
The Adam optimizer with an initial learning rate 0.0002, and a linear decay schedule is used.
The negative loglikelihood loss is used for training the diagnosis and prognosis model.
For prognosis, the parameters of GCN and pooling are initialized using those pretrained on diagnosis task.
All models are trained for 128 epochs with early stopping applied.
4. Q UANTITATIVE R ESULTS
Diagnosis and prognosis performance.
The performance of our model and GCN model with
ASAP are summarized in Table 1.
The highest value, average, and standard deviation of test accuracies on 20 random seeds are reported.
We observed that the Inception V3 feature extraction method constantly outperforms the wavelet decomposition method under the same model conﬁguration.
The gradient explosion occurs in around 50% of the runs under the ASAP model, resulting in optimization failures.
ASAP is very unstable unfortunately, while this issue has not been witnessed during the training of our method.
Thus, the standard deviation of ASAP is much higher.
The training curves using DAP versus ASAP over 20 runs, with varied random seeds and different train-validation-test split are presented in Figure 4.
The ﬁgure shows that our model improves the training convergence, and DAP consistently outperforms ASAP across almost all runs.
Weakly supervised lesion localization results.
Besides graph classiﬁcation, we also use our model to localize the most informative CT slices for each CT scan in the test dataset using the procedure in
Section 3.3 with k s = 10 .
Since the chosen CT slices are the most decisive ones for the diagnosis, they may contain lesions related to COVID-19 or common pneumonia.
Then, we compare our selected slices to the CT slices with lesions labeled by the dataset for each CT scan.
Considering that the CT slices in the same CT scan are sequential, we choose the slices between the ﬁrst and last slices localized by our model or labeled with lesions in the dataset.
The localization of a COVIDpatient is visualized in ﬁgure 5.
Apart from the qualitative visualization, several examples are listed
Method
[Table] [Table]
[Figure] [Figure]
[Figure]
Method
[Table]
in Table 2.
The results are evaluated by Precision, Recall
[Table]
and the Intersection over Union (IoU) for each patient, where IoU =
[Figure]
Lesion ] #[ Localized Slices
[Figure]
w/ Lesion ] .
Through hyperparameter searching, the best performance is achieved when the pooling proportion k is 0.8.
The average IoU of test patients is 41.75% with standard deviation 2.59%.
D ISCUSSION
In this work, we introduce a framework that can provide accurate graph classiﬁcation based on chest
CT scans.
Unlike previous diagnosis methods based on CT scans, our model can produce coarse localization highlighting the potential slices with lesions.
To build trust in this framework and move towards clinical use, we should ensure that the model can explain the reason for their prediction instead of merely outputting the result.
We argue that localization can help analyze prediction failure, and help researchers understand the effect of adversarial attacks in the medical imaging domain.
[Figure]
Supervised Lesion Localization Result of COVID-
Patients
One
[Table]
Localization
[Table]
Slices with
[Table]
Result
[Table] [Table]
Result of
[Table]
Pneumonia Patients
[Table]
As one of the few lesion localization frameworks out there, we argue that the localized slices may help uncover how a given classiﬁcation method arrives at the conclusion.
Speciﬁcally, it might help reveal if a method is merely overﬁtting on the training data by examining the slices it attends to.
This will help build more robust and explainable medical AI models.
Conversely, it may also be used to identify the intrinsic dataset bias, notably the data acquisition bias (Biondetti et al., 2020), thereby guiding the dataset collection process and drive the medical imaging ﬁeld forward.
C ONCLUSION
This paper introduces an efﬁcient and robust GCN-based diagnosis and prognosis system with a distance aware pooling method.
It can cluster nodes and learn the patient-level representation hierarchically.
Apart from accurate diagnosis, a weakly supervised lesion localization method is proposed to support prognosis, making the clinical decision more interpretable and reliable.