MS-ANet: deep learning for automated multi-label thoracic disease detection and classification

Feature extraction networks

The preprocessed input data is fed into two feature extraction networks, where the convolution layers are used as the feature extractor. The employed backbone models for these two feature-extraction sub-networks can be replaced by other models in the proposed flexible multi-scale fusion framework. Unlike the regular CNN, the proposed unified general scheme, namely MS-FIF (Multi-Scale Feature Iterative Fusion), iteratively fuses the feature results extracted from different layers.

The multi-channel attention block MC-AB (Multi-Channel Attention Block), as shown in Fig. 1, has a simple structure and uses two branches with different scales to extract the channel attention weight. One branch extracts the spatial attention of local features, and the other uses Global Avg Pooling averaging to extract the channel attention of global features.

As shown in Fig. 2 above, MS-FIF mainly focuses on the attention problem in the fusion of different scale features in different network structures. Given two feature graphs $X,Y\in R$, $Y$ is a feature graph with a large receptive field range, where MC-AB is a multi-channel attention module, $ZϵR^{C\times H\times W}$ is the fused output feature, $C$ is CONCAT. The large-scale feature maps are down-sampled, so that the two feature maps are the same scale, connected after passing through the multi-channel attention module, and finally output after passing through a multi-channel attention module.

Multi-scale model

We propose a new multi-scale attentional network attention model, which combines the features of different levels or branches to obtain more meaningful features. In Fig. 3, The preprocessed original and rotated input images are fed into two networks with diverse initialization or structure, and fuse the features of four stages through MS-FIF iteration to better fuse semantic and scale inconsistent features. Then, the CAM (Class Activation Map) is used to get the disease attention map, which makes the network only focus on areas with high disease probability. Finally, output the results through the full connection layer.

Figure 3 Multi-Scale Attention Networks. (1) MS-FIF is used to better iteratively fuse features; (2) Using CAM to get the pathogenic attention map, prompting the network to focus only on areas with high pathogenic probability; (3) False frame represents the perceived loss and multi-label balance loss.

$F_A$ and $F_B$ represent the characteristic graphs of networks A and B. Through CAM, a feature map is formed by overlapping the feature set weighted by classifier full connection layer weight $W$, to better make classification decisions based on which feature maps are mainly from the feature set. The losses from two networks are propagated into each other during the backpropagation process, allowing collaborative updating of the networks.

Classifier

As depicted in Fig. 3, the feature maps obtained from the multi-scale attention networks are combined through the global average pooling and fed into a classifier. The proposed method consists of binary classifiers for each disease class. Each binary classifier consists of a fully-connected layer and a sigmoid activation function to compute the probability of being each chest disease.

Perceptual loss

A novel perception loss is defined to extract more significant features in the proposed model. As depicted in Fig. 3, all the features of the last layer of attention feature fusion are obtained in the two networks. Note that the value of each pixel in attention maps indicates different spatial information of the original image. The feature map is then averaged through CAM, and the response size of the feature map is mapped to the original map. The location of different diseases and the size of the lesion area varied greatly on X-ray images and the texture was diverse, and most of them were concentrated in the 'correct' area. A pathogenic attention map is roughly obtained by summing feature maps, computed as follows:

(1) $$S_c=\sum _{x,y}\sum _k{w}_k^c{f}_k(x,y)$$where $S_c$ represents the CAM figure of the $C^{th}$ category, $w\in R^{c\times k}$ represents the weight of the full connection layer, $k$ represents the number of channels, and $f\in R^{C\times H\times W}$ represents the characteristic figure.

For image $x$, $\hat{x}=T\left(x\right)$ is obtained after T transformation, and the corresponding CAM has $\hat{S}=T\left(S\right)$. We obtain two CAMs from the two networks before and after image transformation, $S$ and $\acute{S}$. The L2 loss of two CAM graphs is then computed. The attention loss constrains the cross-attention process smoother by forcing the two networks to find mutually pathogenic regions. The formula for perceived loss is as follows:

(2) $$L_{per}=\sum _i^c{\displaystyle \frac{1}{cHW}{\left|\left|{\hat{S}}_i-{\acute{S}}_i\right|\right|}_2}$$where $c$ represents the number of categories for Label and $S_i$ represents the CAM diagram for the $i^{th}$ category.

Multi-label balance loss

The imbalance between data classes is not conducive for the network to master sufficient texture information; Many samples contain multiple disease information, which is difficult to train; The pathological information of different diseases is different, resulting in different degrees of difficulty in learning. Inspired by Lin et al. (2017), a multi-label balance loss is defined, which is an extension of the focus loss to support multi-classes by adding a balance factor ${\alpha }_i$, extracting more significant features from dominant samples while addressing the imbalance of samples. We aggregate the balance losses for each disease to represent the balance losses for multiple labels:

(3) $$L_{bal}=\sum _{i=0}^c{\alpha }_i{\left(1-\widehat{y_i}\right)}^{p_i}log\left(\widehat{y_i}\right)$$where ${\alpha }_i$ denotes a factor that reduces imbalance between negative and positive samples for i^th class (i = 0…c), with a general value of ${\displaystyle \frac{\left|x_i\right|}{\left|x\right|}}$, where $x_i$ denotes the number of classes $i$. $\widehat{y_i}\left(i=1,2,\cdots ,14\right)$represents the probability value (predictive value) of the network to determine the prevalence of the $i$ disease. Parameter $P_i$ is a difficult-to-easy sample factor. Under certain $P_i$ settings, this may have better performance in mining “difficult” samples, thus making a greater contribution to model training, especially when “easy” samples occupy a large proportion of the data set.

Multi-scale perceptual loss

The multi-scale perceptual loss is defined as the combination of perceptual and multi-label balance losses as follows:

(4) $$L=\alpha L_{per}+L_{bal}$$where $\alpha$ is a coefficient parameter that controls the trade-off between two losses. More weight on L_per induces a more accurate model, focusing on the region where the disease occurs. In contrast, more weight on the multi-label balance loss increases the loss weight of difficult-to-identify diseases and reduces the loss weight of easy-to-identify diseases, thereby enhancing the network's learning of difficult-to-identify samples.

Dataset and training

In this section, the proposed network is evaluated on the Chest X-Ray14 dataset (Wang et al., 2017), which comprises 112,120 frontal-view X-ray images of 30,805 unique patients with the text-mined fourteen disease image labels (where each image can have multi-labels), mined from the associated radiological reports using natural language processing. According to the official ratio, the dataset is divided into training and testing sets to ensure a fair comparison. In the experiment, we zoom the input image to 256 × 256, and randomly select the center point to cut the image to 224 × 224. Random rotation and random flip data enhancement methods are used in the training process. Among training data, 78,485 images and 8,039 images are used for training and validation, respectively. No overlap among the three patient-sets is ensured. The data batch size is 128, and the learning rate is initially set as 0.001. The dropout rate is set as 0.2 for the last fully-connected layer, and α is set to 0.0002. The $P_i$ in the multi-label balance loss is set to 1.5 according to the existing research. All experiments were evaluated based on the AUROC value.

Comparison analysis

Table 1 presents the quantitative comparison of the performance of different models. Multi-Scale Attention Network—MS-ANet₁, MS-ANet₂, and MS-ANet₃ are constructed with different backbone networks. MS-ANet₁ is the multi-scale attention network whose backbone networks are two densenet121 networks. MS-ANet₂ represents the multi-scale attention network model with densenet121 and densenet169 backbone networks. MS-ANet₃ represents the multi-scale attention network model with two densenet169 backbone networks. Each sub-network is differently initialized using distinct weights. Note that additional localization labels were used in the experiments of Li et al. (2018) in addition to the official segmentation data. For a fair comparison of the experiments, the results are not considered for the best results of each row (marked in bold).

As shown in Table 1, the multi-scale attention model has achieved the best results in terms of the average AUROC scores for most individual diseases. The proposed model induced significant improvement, especially for an unbalanced disease class where positive samples are rare. For instance, in the dataset, 'Hernia' has only 227 pictures (0.202 %), and ‘Emphysema’ has only 2,516 pictures (2.244 %). The proposed MS-ANet₂ model obtains outstanding AUROC values, 0.947 and 0.934, respectively, for Hernia and Emphysema, outperforming other methods. It is mainly due to the employed multi-scale attention loss, punishing those samples that are difficult to identify so that negative and positive samples can be distinguished better.

Ablation experiment

In order to evaluate the contribution of each module of the proposed model, ablation experiments were conducted by comparing the impact on classification accuracy. In Table 2, Densenet121 and Densenet169 are denoted as ‘121’ and ‘169’, respectively. The binary cross-entropy, multi-label balance, and perceptual losses are denoted as L_bce, L_bal, and L_per, respectively.

Table 2 shows that the proposed MS-ANet improves the accuracy in the unbalanced dataset. The AUROC score is increased by 0.032 on average with the MS-FIF, obtaining 0.840 from 0.806. Since through a more collaborative way through each other's s reverse propagation gradient to update. The AUROC score is increased to 0.847 with adopting the multi-scale perceptual loss. Since the perceptual loss narrows the attention of the two networks, thus promoting the cross-attention. The proposed MS-ANet obtains 0.847, 0.850, and 0.845 AUROC, which is the latest and most advanced result.

Image-level supervised disease localization

Based on the quantitative analysis of the classification performance of the model, this paper makes a qualitative analysis of the model. The localization heat map of the lesion area is generated by the CAM, which is used to visually explain the lesion area on which the network is based when judging the disease. The product of weight and feature mapping of the fully-connected layer and pooling layer is directly located because of the global average pooling used as the final pooling layer.

Figure 4 presents the localization results of eight diseases, in which the bright red part is the main area for diagnosis, and the blue boundary frame indicates the reallocation of the disease. After observation, whether it is a large disease in the lesion area (such as Fig. 4B Cardiomegaly) or a small disease in the lesion area (such as Fig. 4F Nodule), the lesion area located by the thermal map can be well consistent with the lesion area annotation box given by the doctor, which further verifies that the feature information based on the network in the diagnosis is accurate and effective. Besides, visualization of lesion areas can provide better visual support for professional doctors by computer-aided diagnosis in clinical application, obtain doctors' trust, and help doctors to make a rapid and accurate diagnosis.

Dataset and training

The CheXpert (Irvin et al., 2019) dataset, which contains large-scale data of 224,316 chest radiographs of 65,240 patients with frontal and lateral radiographs, is also used to validate the proposed model. The experiment was conducted as similar to the experiment on the Chest X-Ray14 dataset in all comparative experiments. The images are divided into 222,914 images for the training set and 734 images for the testing set. No duplication is ensured between those images.

Comparison analysis

In this experiment, the U(uncertain)-One introduced in CheXpert’ s paper (Irvin et al., 2019) is only used to maintain consistent experimental settings, labeling all uncertain samples to 1. Besides, we tested the models for fourteen labels and divided them into using only the frontal radiographs, and using the frontal and lateral radiographs equally. The classification results for fourteen classes are compared and summarized in Table 3. As shown in Table 3, the proposed model obtains the highest AUROC scores (marked in bold) for most classes, including both the frontal radiographs and the frontal/lateral radiographs.