Adaptive Image-Feature Learning for Disease Classification Using Inductive Graph Networks

TL;DR

This paper introduces an end-to-end inductive graph neural network architecture for disease classification that jointly trains image and graph filters, leading to improved accuracy and stability over existing methods.

Contribution

The proposed architecture enables joint training of CNN and graph filters in an end-to-end manner, enhancing disease classification performance.

Findings

Significantly improved classification scores on a modified MNIST dataset.

Achieved comparable results with higher stability on chest X-ray images.

Demonstrated the impact of graph structure on feature learning.

Abstract

Recently, Geometric Deep Learning (GDL) has been introduced as a novel and versatile framework for computer-aided disease classification. GDL uses patient meta-information such as age and gender to model patient cohort relations in a graph structure. Concepts from graph signal processing are leveraged to learn the optimal mapping of multi-modal features, e.g. from images to disease classes. Related studies so far have considered image features that are extracted in a pre-processing step. We hypothesize that such an approach prevents the network from optimizing feature representations towards achieving the best performance in the graph network. We propose a new network architecture that exploits an inductive end-to-end learning approach for disease classification, where filters from both the CNN and the graph are trained jointly. We validate this architecture against state-of-the-art inductive graph networks and demonstrate significantly improved classification scores on a modified MNIST toy dataset, as well as comparable classification results with higher stability on a chest X-ray image dataset. Additionally, we explain how the structural information of the graph affects both the image filters and the feature learning.

Tables

Table 1. Table 1: Performance on MNIST dataset. Aff. shows which affinity mechanism was used: θ 𝜃 \theta = l1-distance, L=labels, CL=labels with connection between 5 and 6 . P-val is reported against CNNGAT. Occ. shows average correct class probability of softmax function for occlusion shift. O. p-val shows p-value compared to CNNGAT.

Network	Aff.	Accuracy	Lowest class acc	p-val	Occ.	O. p-val
CNN	/	0.822 $\pm$ 0.005	0.69 $\pm$ 0.049	10e-5	0.69	0.0
RawGAT	$\theta$	0.794 $\pm$ 0.003	0.450 $\pm$ 0.008	10e-5	/	/
SkipGAT	$\theta$	0.829 $\pm$ 0.014	0.538 $\pm$ 0.066	10e-5	0.68	0.0
EndGAT	$\theta$	0.881 $\pm$ 0.016	0.780 $\pm$ 0.087	10e-4	/	/
CNNGAT	$\theta$	0.911 $\pm$ 0.005	0.810 $\pm$ 0.035	/	0.81	/
CNNGAT	Rand.	0.816 $\pm$ 0.004	0.613 $\pm$ 0.052	/
CNNGAT	L	0.980 $\pm$ 0.036	0.956 $\pm$ 0.079	/
CNNGAT	CL	0.992 $\pm$ 0.001	0.989 $\pm$ 0.003	/

Table 2. Table 2: Performance of CNN and CNNGAT on the described NIH ChestX-ray 14 dataset for 16,000 images. The column setting is identical to Tab. 1 .

Network	Aff.	Accuracy	Lowest class acc	p-val	Occ.	O. p-val
CNN	/	0.429 $\pm$ 0.015	0.144 $\pm$ 0.102	0.109	0.27	2.6e-5
CNNGAT	meta	0.437 $\pm$ 0.014	0.086 $\pm$ 0.070	/	0.30	/

Equations

$${\alpha }_{ij}=\frac{\exp (\sigma ({\vec{a}}^T([\text{W}{\vec{v}}_i||\text{W}{\vec{v}}_j])))}{{\sum }_{r\in N_i}\exp (\sigma ({\vec{a}}^T[\text{W}{\vec{v}}_i||\text{W}{\vec{v}}_r]))}$$

$${\overrightarrow{v^{\prime }}}_i={\Vert }_{k=1}^K\sigma \left(\sum _{j\in N_i}{\alpha }_{ij}^k{\text{W}}^k{\vec{v}}_j\right),$$

Full text

1 Introduction

In the past few years, the potential of data processing with Geometric Deep Learning (GDL) on non-Euclidian data structures has been demonstrated by a large body of work [1, 2, 6, 7, 8]. One uprising application, which is also the subject of this work, lies in the field of computer-aided disease classification (CAD), where the usage of patient meta-information is leveraged to build a graph system of meaningful patient inter-relations [8, 5]. These patient relations allow e.g. an improved classification of image data, which enables the prediction of a disease. This is achieved by using various forms of graph signal processing for the training of localized graph filters, where patients are modeled as vertices in the graph, and images as feature vectors. One family of approaches [1, 6, 8] are graph spectral methods, which utilize the normalized graph Laplacian $L$ to perform a Fourier transform in order to find optimal filters in the frequency domain. A major drawback of spectral methods is the limitation of their learned representation to the graph they are trained on. Therefore, their transductive approach requires a retraining for every new graph system, making them difficult to use for a direct application on newly received patient data. We therefore focus on non-spectral methods, which overcome this limitation by working on local graph structures, mimicking the standard convolution operation on the spatial relations of a graph node. The concept relevant for this work are composition-based spatial methods that use multiple consecutive layers to update the vertex representation. Hamilton et al. presented GraphSage [3], which samples a fixed number of neighbors of a vertex and aggregates them to receive an updated representation of the feature vector in every layer. This inductive approach allows the successive processing of every vertex of the graph. The system however did not define a preference of which neighboring feature representations are most relevant for the update. We thus focus on and adapt Graph Attention Networks (GAT) in this work. GAT was introduced by [10] and leverages an attention mechanism to learn this importance measure in order to optimize the representation learning.

Regarding CAD systems, in [8] a GDL-approach for disease classification was proposed, modeling patient similarities through meta-information and performing patient classification in the ABIDE and ADNI datasets using Graph Convolutional Networks (GCN). In 2018, self-attention was used in [5] to optimize the usage of meta-information in the TADPOLE dataset. Both methodologies used pre-learned features in their approach. Since we are evaluating our method on chest X-ray images, we are also mentioning related works out of this field. In 2017, Wang et al. [11] proposed a new chest X-ray database, known as “ChestX-ray8”, which contains 108,948 frontal view X-ray images of 32,717 unique patients for eight different disease classifications and performed multi-label disease classification and localisation using class activation mapping [12]. In [4], disease classification was realized on three chest X-ray datasets using standard deep convolutional networks, localizing activations with occlusion mapping. CheXNet [9], a DenseNet of 121-layers, was used for classification of all 14 diseases for the NIH ChestX-ray 14 dataset, with an emphasis on pneumonia. All the works relied on standard CNN approaches and did not leverage any non-imaging information.

Contributions. Although GDL approaches used for CAD showed promising performance, two major drawbacks are present in the methodologies developed so far. First, it is necessary to extract the features in a pre-processing step, which results in additional effort. Second, a potentially not optimized feature representation for the usage inside the graph system is learned. To address these limitations we propose CNNGAT, a novel method that yields an inductive end-to-end approach for both learning an optimized feature representation and performing the graph convolutions during training time. Additionally, the introduced concept leverages inter-class connections to specifically improve the performance on instances from classes which are more diffcult to distinguish. To the best of our knowledge this is the first work that combines the training of the CNN and the graph network to obtain optimized feature representations. We evaluate our proposed method on the ChestX-ray 14 [9] and a modified MNIST dataset.

2 Methodology

2.1 Explanation of the Proposed Model

General Framework. Our proposed network CNNGAT performs the classification of images represented as 2D intensity arrays X using non-imaging meta-information in an inductive end-to-end approach. It tries to optimize the objective function $f(\textbf{X},G(\textbf{V},\textbf{E})):\textbf{X}\rightarrow\textbf{Y}$, where $G(\textbf{V},\textbf{E})$ is a graph with vertices V containing a feature representation of every image $\vec{x}_{i}$ and binary edges E as connections between the vertices and Y as set of classes. The whole set of feature representations for every vertex $\vec{v}_{i}$ of the graph is therefore defined as: $\textbf{V}=\{\vec{v}_{1},\vec{v}_{2},...,\vec{v}_{N}\},\vec{v}_{i}\in\mathbb{R}^{F}$, where $F$ is the dimension of the representation. The edges E are created based on non-imaging information belonging to every image $\vec{x}_{i}$. The CNN extracts a feature representation of image $\vec{x}_{i}$ and distributes it to the corresponding vertex $\vec{v}_{i}$. Subsequently, the obtained representation is processed by two consecutive GAT layers. The resulting representation is concatenated with the initial one of the CNN to perform the final prediction (Fig. 1).

Neighborhood Concept. A GAT layer updates the feature representation $\vec{v}_{i}$ based on the ones of neighboring vertices. Therefore, it requires that a neighborhood of images is processed together with $\vec{x}_{i}$. Since the general size of image datasets prohibits the use of all neighbors for a vertex $\vec{v}_{i}\in\textbf{V}$, different from [10], only a fixed sub-sample of size $n$ of all vertices $\vec{v}_{j}$ with $e_{ij}\in\textbf{E}$ is considered as neighborhood $N_{i}$ of $\vec{v}_{i}$. If for a vertex $\vec{v}_{i}$, the size $n$ of the neighborhood cannot be reached, a random selection of $N_{i}$ is used multiple times until the amount $n$ of vertices is achieved. As described in [3], this procedure has the advantage of setting the computational effort of the network to a fixed amount.

GAT Layer. To achieve a higher representation of $\vec{v}_{i}\in\mathbb{R}^{F}$ with new dimension $F^{\prime}$, a shared learnable linear transformation $\textbf{W}\in\mathbb{R}^{F^{\prime}\times F}$ is applied to $\vec{v}_{i}$ and its considered neighbors $\vec{v}_{j}\in N_{i}$. Then, a shared attention mechanism $a$ that consists of a single-layer feed-forward network is used to obtain the attention coefficient $\alpha$, representing how important $\vec{v}_{j}$ is for the update of $\vec{v}_{i}$ and computed as $a(\textbf{W}\vec{v}_{i},\textbf{W}\vec{v}_{j})=\vec{a}^{T}[\textbf{W}\vec{v}_{i}||\textbf{W}\vec{v}_{j}]$, with the concatenation (symbolized as $[~{}||~{}]$) of $\textbf{W}\vec{v}_{i}$ and $\textbf{W}\vec{v}_{j}$ and $\vec{a}\in\mathbb{R}^{2F^{\prime}}$. After activating the calculated attention using the activation $\sigma=$ LeakyReLU, the softmax function is applied for all $\vec{v}_{j}\in N_{i}$ to obtain normalized and easily comparable attention coefficients $\alpha$:

[TABLE]

Every attention coefficient is multiplied with the feature representation $\textbf{W}\vec{v}_{j}$ and added up to obtain the new representation $\vec{v^{\prime}}_{i}$. To stabilize the representation, this procedure is repeated multiple times with individual $\textbf{W}^{k}$, called heads, performing the same attention mechanism for a vertex $\vec{v}_{i}$. The representations $\vec{v^{\prime}}_{i}$ are concatenated (represented as $\|$), yielding the new feature representation:

[TABLE]

where $\alpha_{ij}^{k}$ is the attention coefficient of head $k$ for the vertices $\vec{v}_{i}$ and $\vec{v}_{j}$, and $K$ is the number of heads [10].

CNN/GAT Hybrid. The initial representation $\vec{v}_{i}$ is processed at the same time in a skip connection using a single-layer feed-forward network $\textbf{W}_{skip}$. The output of the layer is activated using LeakyReLU and then concatenated with the output of the graph attention network $\vec{v^{\prime}}_{i}$ to receive the hybrid representation $\vec{v^{\prime\prime}}_{i}=[\vec{v}_{i}||\vec{v^{\prime}}_{i}]$. This strategy enables both the initial representation of the CNN and the altered representation after neighborhood interaction to be used for the prediction, leveraging both information. After a final activation, a last single-layer feed-forward network is used to map the concatenated vector to the class output and apply a softmax function. The gradients are backpropagated through the whole network, updating not only the graph network but also the CNN, and therefore updating the feature extraction in every iteration.

2.2 Motivation of Skip Connection for CNN/GAT Hybrid

Two concepts motivate the concatenation of the representations extracted from the CNN and the GAT. First, the skip connection allows a direct gradient propagation to the CNN, thus fortifying a proper filter learning. More importantly, the combination of both outputs enables a comparison of the pure feature extraction of the CNN and the “impure” feature representation after its interaction with the neighborhood (Fig. 1 top right). The individual connectivity of instances from a class to other classes in the graph can yield an additional unique contribution to the feature representation after the feature aggregation process of GAT, if these inter-class connections are different for every class. Especially for classes usually difficult to distinguish, this individual contribution of feature representations from other instances significantly improves the classification performance on these classes since their originally similar feature representations are altered and more distinct compared to their initial representation transported by the skip-connection. Therefore, the interesting and unintuitive observation is obtained that inter-class connections can be highly beneficiary in the used setup.

3 Experiments

3.1 Datasets

First, a modified version of the MNIST dataset is used to prove the concept of the proposed method. MNIST consists of 70,000 handwritten digits from [math] to $9$. Only the lower half of every image is used as image input, while the upper half is processed for the creation of the affinity graph. The objective is to show the improved classification of classes difficult to distinguish like 3 and 5 due to their identical lower shape. Following the inter-class setup explained in Sec. 2.2, we additionally take into account the digit 6, which shows a similarity to only the number 5 in the upper half. We therefore create a subset of 6,000 images of the numbers 3, 5 and 6 for training and 2,860 images for testing.

Secondly, we adapt the approach to the task of disease classification on the NIH ChestX-ray14 dataset. The dataset consists of 112,120 labelled X-ray images, containing one or more of 14 different diseases, with substantial imbalance for some classes. 16,000 single-label images of the eight diseases suggested by [11] are taken into consideration to obtain a potentially clearer patient cohort system.

3.2 Experimental Setup

Affinity Graph Construction. For MNIST, the construction of the affinity graph is based on the upper half of the image. As a metric, the $l1$-distance of the flattened image intensity vectors is used. If the distance between the two image vectors $v_{1}$ and $v_{2}$ is below the threshold $\theta$, the edge is established. The best threshold was found with $\theta=0.1$ for the average pixel intensity difference on an intensity scale from 0 to 1. Then, to analyze the proposed network and to prove that the obtained improvement is related to the previously described meaningful inter-class connections, three additional settings are used: 1) random edges, 2) edges based on label information, and 3) edges based on labels including the helpful inter-class connection between 5 and 6. It has to be stated that the usage of label information is just used here to prove the concept of the network.

For the NIH chest X-ray dataset, the accessible meta-information is used to create the affinity graph. Here, we are following the current state of the art introduced by [8] to leverage an age and gender difference metric. After performing an analysis of the statistics of the age and gender distribution with respect to the patients’ diseases, we connect subjects in our experiments if at least one of the following three conditions is fulfilled for two patients $\textsc{p}_{1}$ and $\textsc{p}_{2}$:

1. patient id${}_{\textsc{p}_{1}}$ = patient id${}_{\textsc{p}_{2}}$

2. gender${}_{\textsc{p}_{1}}$ = gender${}_{\textsc{p}_{2}}$ and $|$age${}_{\textsc{p}_{1}}$ - age${}_{\textsc{p}_{2}}|\leq 1y$

3. gender${}_{\textsc{p}_{1}}$ $\neq$ gender${}_{\textsc{p}_{2}}$ and age${}_{\textsc{p}_{1}}$ = age${}_{\textsc{p}_{2}}$

Network Setup. We create three baselines and one alternative method to analyze the different aspects of CNNGAT. To show the improved learning of feature representations we compare against the performance of a simple CNN, the normal GAT network on pre-learned features by the CNN (RawGAT), and our CNNGAT with pre-learnt features instead of end-to-end training (SkipGAT). To validate the hybrid approach we compare against our end-to-end CNNGAT network without using skip connections (EndGAT).

Feature Extraction. For MNIST, the CNNGAT was trained with a simple CNN containing two convolutional layers (32 and 64 channels). For the NIH dataset, AlexNet was used to easily prevent overfitting on the 16,000 images.

Parameters. The following parameters were chosen:

MNIST: 60 features as GAT input, 2 GAT layers (30 and 10 units), 5 heads, dropout: 0.3, neighbors: 4, weight decay: 5e-3, lr: 0.02, lr $\times$ 0.3 every 20 epochs.

NIH: 60 features as GAT input, 2 GAT layers (30 and 32 units), 5 heads, dropout: 0.3, neighbors: 4, weight decay: 5e-3, lr: 0.01, lr $\times$ 0.3 every 30 epochs.

3.3 Modified MNIST Dataset

The results on the above described MNIST toy dataset shown in Tab. 1 indicate that the CNNGAT is indeed improving the classification by using the inter-class connections for numbers 5 and 6. Using pre-trained features in the proposed network approach (SkipGAT) diminishes the performance of the system, showing the value of adaptive feature learning. Also the approach without Skip-Connection (EndGAT) is significantly outperformed (Wilcoxon signed-rank test, significance level $p<0.05$). To give further evidence that the features are learned in a more robust way, we perform a shifted occlusion for 1000 random images. A occluding window is slid across the image and after prediction is performed, the probability of the correct class in the softmax function is recorded for every position. One example is shown in Fig. 2 a-d) for an instance of the number 3. It is visible that the classification is significantly more stable compared to the other networks, also stated in Tab. 1.

For the second setting of graph edges described in Sec. 3.2, the lower part of Tab. 1 clearly shows that the presence of meaningful inter-class connections is leading to top-level performance even outperforming a clean label graph.

3.4 NIH ChestX-ray 14 Dataset

Tab. 2 shows the results on the ChestX-ray 14 dataset. The CNNGAT had a slightly better performance than the used CNN, but the improvement is not significant. The absence of the effects seen on the MNIST dataset are addressed in the discussion. The occlusion shift analysis however shows that the classification is stabilized by the graph (Fig. 2 right). The probability of predicting the correct class under occlusion is significantly higher compared to the raw CNN (Tab. 2).

4 Discussion and Conclusion

The results on the MNIST dataset follow the expectation of the network performance. Only when training the CNN and GAT end-to-end, optimized performance is reached. Especially the comparison to SkipGAT, which is identical to CNNGAT except its pre-learned features, states the importance of the feature learning during training time. Additionally, the fact that the EndGAT setting is significantly outperformed fortifies the usage of the CNN/GAT hybrid approach. This behaviour was not clearly reproducible on the ChestX-ray 14 dataset. After performing a statistical analysis on the data distributions we hypothesize that the available meta-information was not relevant enough to build the required patient cohorts. Generating an artificial meaningful graph on the dataset showed the expected clear performance improvement.

We have proposed a new methodology to train a CNN and GAT network end-to-end and used their output in a hybrid approach to leverage both advanced feature learning and inter-class feature representations. The experiments clearly showed the superiority of the approach under the constraint that a meaningful graph can be created. In future work, the network architecture can be further optimized by e.g. including a learning mechanism for the adjacency system. With an automatized graph creation the necessity of finding meaningful meta-information is erased what makes the network applicable in a more general way.

Acknowledgements. The study was supported by the Carl Zeiss Meditec AG in Oberkochen, Germany, and the German Federal Ministry of Education and Research (BMBF) in connection with the foundation of the German Center for Vertigo and Balance Disorders (DSGZ) (grant number 01 EO 0901). Further, we thank NVIDIA Corporation for the sponsoring of a Titan V GPU.