BEYOND COVID-19 DIAGNOSIS: PROGNOSIS WITH HIERARCHICAL GRAPH REPRESENTATION LEARNING Anonymous authors Paper under double-blind review

Abstract

Coronavirus disease 2019 , the pandemic that is spreading fast globally, has caused over 34 million confirmed 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 convolutional networks (GCNs). The chest CT scan of a patient, typically involving hundreds of sectional images in sequential order, is formulated as a densely connected weighted graph. A novel distance aware pooling is proposed to abstract the node information hierarchically, which is robust and efficient 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 94.8% accuracy for diagnosis. It has a 2.4% improvement compared with the baseline model on the same dataset. In addition, we can localize the most informative slices with disease lesions for COVID-19 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 efficiently and better support clinicians to determine the proper clinical treatment.

1. INTRODUCTION

Coronavirus disease 2019 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 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) . 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 Veličković 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



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. (

