Learning Hierarchical Attention for Weakly-supervised Chest X-Ray Abnormality Localization and Diagnosis

As noted above, we generate two diagnostic attention maps in attention hierarchy: the positive attention map $M^{p}$ with shape $1\times H\times W$ and $D$-way abnormality attention maps $M^{a_{k}}$ with shape $D\times H\times W$ (where $D$ is the number of abnormality classes and $k=1,2,\ldots,D$ indicates the specific abnormality). Given the hierarchical relationship between these two attention maps discussed above, we seek a learning objective that enables the model to learn to produce attention that respects this hierarchy. Specifically, given a CXR image, all abnormality attention maps $M^{a_{k}}$ shall be contained within the region bounded by the positive attention map $M^{p}$ as shown in Fig. 4. Accordingly, our proposed attention bound objective, $L_{bound}$, attempts to spatially constrain the region covered by each $M^{a_{k}}$ with respect to $M^{p}$, and is realized for input CXR image as:

where the image’s ground-truth label $y_{k}\in\{0,1\}$ ($0/1$ indicates absence/presence of abnormality $k$), $N$ is the number of positive classes in the image label set $\{y_{k}\}$, $i$ and $j$ represent the $(i,j)^{th}$ pixel in the corresponding attention map, and $T(M_{ij}^{{a_{k}}})$ is the soft masking operation defined in Equation 2 which masks out the impact of background noise in the attention maps. In summary, the $L_{bound}$ can ensure abnormality attentions $M^{a}$ to lie within the positive attention $M_{p}$.

While our proposed $L_{bound}$ helps enforce explicit spatial constraints on the spatial bounds of the abnormality attention $M^{a_{k}}$, it does not, however, provide any direct supervision on the spatial extent of the positive attention $M^{p}$. To ensure the positive attention only encompasses all the possible abnormality attention regions, and no more than that, we enforce explicit union constraint on the spatial extent of positive attention map $M^{p}$ and the union of all the abnormality attention maps $M^{a_{k}}$. Specifically, as shown in Fig. 4, we seek the spatial extent of the positive attention map $M^{p}$ to be no more than the spatial extent of the union of all the abnormality attention maps $M^{a_{k}}$. To this end, we first compute the union of all abnormality attention maps, denoted $M^{u}$, as:

where similar to $L_{bound}$, we only consider the positive abnormalities. Then, our proposed attention union loss, $L_{union}$, is realized by calculating the region overlap between $M^{p}$ and $M^{u}$ as:

where $T(M_{ij}^{p})$ refers to the same soft-masking operation of Equation 2.

We conduct experiments on the NIH Chest-Xray14 dataset [8] and CheXpert dataset [9], two large-scale CXR datasets. In our experiments, we use dilated ResNet50 [33] as the backbone. The size of the abstracted feature map to generate the various attention maps is $\frac{1}{8}$ of the input images. We use the Adam [34] optimizer with momentum set to $0.9$, a weight decay of $0.0001$, and a learning rate of $0.00002$ that is reduced by a factor of 10 after every $10$ epochs. We train our model with a batch size of 24 and empirically set $r=6$ for the LSE pooling, as defined in [32]. During training, we perform random translation and horizontal flipping for augmentation. We resize the original 3-channel images to $512\times 512$ as the input. The loss weight factors $z$, $\lambda_{1}$, $\lambda_{2}$, and $\lambda_{3}$ are set to 0.5, 0.01, 0.001, 0.001, respectively.

We use the area under the ROC curve (AUC) to measure classification performance, and intersection over union ratio (IoU) and the intersection over the detected region (IoR) to quantify localization results. To calculate IoU and IoR, we use bounding boxes of localized attentive regions. Following prior work [8, 12], we report the ratio of the number of cases with correct localization against the total number of cases in each class. A localization result is considered correct if the criterion of either IoU $>T(IoU)$ or IoR $>T(IoR)$, where $T(*)$ is the threshold, is met.

We conduct two experiments to evaluate the performance on the abnormality classification task and compare to the state-of-the-art (SOTA) methods. The first experiment is based on the official split [8] and prepared at the patient-level. All images with box annotations are in the testing set and not used during model training. Results are shown in Table 1, where our method outperforms state-of-the-art methods in terms of AUC. In particular, our method outperforms DNetLoc [3] and CAN [10] methods that employ deeper models (DenseNet121) as their backbones.

In the second experiment, we use the 5-fold cross-validation (CV) scheme following [12, 21]. For the convenience of comparison, we use the same notations with [12, 21] by referring to unannotated CXR images as those with only image-level labels and annotated CXR images as those with both image- and box-level labels. In each fold, we use $70\%$ of the annotated and $70\%$ of unannotated images for training, and $10\%$ of the annotated and unannotated images for validation. Then, the rest $20\%$ of the annotated and unannotated images are used for testing. We summarize AUC scores of the compared methods w.r.t. the 14 abnormalities of testing set in Table 2, where our method obtains the best mean AUC score of $0.835$.

Two experiments for assessing localization performance are conducted. In the first experiment, we show localization performance by considering only image-level annotations, i.e., abnormality labels. Following [12, 21], we train our model with $100\%$ images (111,240) without any box annotations and test with the 880 images with box annotations. We summarize our results in Table 3, where we compare to SOTA methods at $T(IoU)=0.1$ and $0.3$. Our method significantly improves localization performance for all abnormalities at $T(IoU)=0.1$, with a mean score of $0.75$. In particular, we achieve more than $150\%$ improvement in localization for “Pneumonia” compared to the results of the most recent SOTA method [21]. We also observe remarkable localization improvements for the “Atelectasis”, “Mass” and “Nodule” abnormalities. In $T(IoU)=0.3$, our method also yields competitive localization performance and outperforms [21] for “Mass”, “Nodule” and “Pneumonia”. We further show qualitative comparison of the attention maps with SOTA method [21] in Fig. 6, where the exactness of our attention map can be corroborated.

Also, we notice that our proposed method does not perform well for “Cardiomegaly”, “Effusion” and “Infiltration” at $T(IoU)=0.3$ compared with Liu et al. [21]. The underlying reasons may be twofold: 1) the threshold setting for the attention maps; 2) the coverage range of bounding boxes. For the first reason, we set a very high threshold (0.999) on attention maps to get the final binary localization masks for the computation of evaluation scores in Tables 3. This threshold is very rigorous and may generate relatively small object masks, which may not be favorable for large class like “Cardiomegaly”, for localization. For comparison, we further explore a smaller attention threshold of 0.1. The resulting correct ratio score at $T(IoU)=0.3$ with the attention threshold of 0.1 for “Cardiomegaly” is 0.72, which is slightly higher than result of Liu et al. [21]. The second reason is the issue of box definition that may lead to lower performance of our method at $T(IoU)=0.3$, particularly for “Effusion” and “Infiltration”. We have invited an experienced radiologist to carefully review the original annotated boxes and suggest that the variety of the annotation coverage for some classes (e.g., “Effusion”, “Infiltration” and “Pneumonia”) is very large (see figures and detailed analysis for better illustration in the supplementary material). Larger boxes with more non-related regions may favor localization results with relatively larger masks for higher scores. Specifically, the results of Liu et al. [21] for the “Infiltration” and “Effusion” cases in the Fig. 6 are fuzzy and cover more non-related regions (even the abdominal region), while our localization results are closer to the abnormal regions of “Effusion” and “Infiltration”. Therefore, the results of Liu et al. [21] for “Effusion” and “Infiltration” may be higher.

In the second experiment, we use the 5-fold CV scheme following [12, 21]. In each fold, we train our model with $50\%$ of the unannotated images and $80\%$ of the annotated images, and tested with the remaining $20\%$ of the annotated images. As shown in Table 4, our method outperforms SOTA methods in most cases. In particular, for the difficult classes of “Mass” and “Nodule”, our localization performance is remarkably better. Similar with results in Table 4, our method does not perform well for some abnormalities (e.g., “Effusion” and “Infiltration”). Like previous observation, it may be due to large box annotations for these abnormalities. For “Pneumonia”, we notice that our method achieves much better performance than baselines at $T(IoU)<0.5$ but much worse at $T(IoU)>=0.5$. Although most boxes provided by NIH for “Pneumonia” are reasonable, some box annotations also enclose many non-related regions. Meanwhile, it is worth noting that there are some missing boxes in NIH Chest X-ray14. The figures and more details can be found in the supplementary material. Thus, we suggest the results with IoU scores in the range of 0.3 to 0.7 are acceptable (see Fig. 9 and Fig. 10 in the supplementary material). It can be observed that nearly 63% ($0.69-0.06$) of our results have IoU scores of “Pneumonia” in the range of 0.3 to 0.7, compared to 24% ($0.49-0.25$) of results in Liu et al. [21].

To better illustrate the localization performance, we also show the $T(IoR)$ results in Table 5. It calculates the intersection over the detected region, reducing the effect from non-related regions in the box annotations in evaluation. As shown in Table 5, our method gives the best performance in most items of $T(IoR)$ values. At the same time, we show the performance of the self-refinement method which can reduce the noise of the foreground in the box annotations during training. It proves that the self-refinement method can make the attention results more concentrated within the box annotations. We show qualitative results of the self-refinement method in Fig. 7.

In summary, our proposed hierarchical attention framework can result in more precise and pathological plausible localization of abnormalities and achieves state-of-the-art results at several $T(IoU)$ and $T(IoR)$ settings in our experiments. It is worth noting that our method localizes the relatively small abnormalities like “Nodule” much better. The detection of pulmonary “Nodule” in CXR images is also a very challenging task with low sensitivity in the range of $0.69$ to $0.82$ for radiologists [35].

Table 6 shows results of the ablation studies with the AUC metrics and localization results of $T(IoU)=0.1$, using the official split for training and testing. Abnormality attention is applied in all the models to generate the abnormality localization prediction, which is the same as CAM [13] when not using our attention bound and union loss.

First, ResNet50 without the positive/negative classification branch is implemented with down-sampled input images (factor of $32$). The second backbone is the dilated ResNet50 with images down-sampled by a factor of $8$. The global average pooling (GAP) is utilized for the two backbones. The corresponding results are shown in the first two rows of “RN50-32-GAP” and “RN50-8-GAP”, respectively, in the Table 6. We can see that higher resolution, i.e., smaller down-sampling factor, is helpful for the localization of classes like “Mass” and “Nodule”. Second, we propose to use the LSE pooling instead of GAP. To illustrate the effectiveness, we compare the network of “RN50-8-LSE” to “RN50-8-GAP”. As shown in Table 6, the average localization score is higher with LSE.

Third, we perform an ablation study for different versions of FAB. Specifically, we compare three attention configurations: FAB-C (channel-wise only), FAB-P (position-wise only), and FAB (both channel-wise and position-wise). The combination of both attention modules achieves the best average localization performance (0.61).

Fourth, we add the positive/negative classification branch to generate the positive attention (“RN50-8-LSE+FAB+PA”), with which the average localization score (0.63) is slightly higher than “RN50-8-LSE+FAB”. However, with the addition of our attention bound and union loss in the last row, the average localization performance increases to 0.70. The results suggest that the proposed attention losses are important for the hierarchical attention representation in guiding the learning of abnormality attention from positive attention. Results of our method without FAB (“RN50-8-LSE+PA+ABU”) suggest that the FAB module is efficient for refining the feature encoding from the backbone network. Comparing the results of “RN50-8-LSE+FAB” and the last row, we can see that the localization results of the “Infiltration” and “Pneumonia” are significantly boosted by our hierarchical attention mining method. In particular, the “Pneumonia” class has the smallest number of samples (876 images) in the training set (86,524 images) of the official split. It is evident that our HAM method can alleviate the data-imbalance problem even in such an extreme situation.

At the same time, we show the results of our method with only attention bound constraint (“RN50-8-LSE+FAB+PA+Bound”). It performs worse in both classification and localization tasks than the model with attention bound and union losses. Because there is no constraint for positive attention without the incorporation of attention union loss, the incorrect prediction of positive attention may misdirect the abnormality attention. In such a case, the final classification and localization performances are compromised. Accordingly, by comparing the ablation of attention bound and attention union losses, the effectiveness of the synergy of two attention losses is corroborated.

To better illustrate the statistical significance of our method, we also calculate the $p$-value between the baseline “RN50-8-LSE” and our model “RN50-8-LSE+FAB+PA+ABU” in Table 6. The $p$-value for the classification predictions of these two models is $4.31\times 10^{-5}$, implying that the proposed methods have significant improvements compared with “RN50-8-LSE”.

Finally, the effectiveness of the proposed AMSE loss is shown by comparing the performance to the original MSE loss. We conduct the same 5-fold CV experiment as the second experiment in section 4.2.2. As shown in Table 7, the AMSE loss is especially effective for the smaller abnormalities (e.g., “Nodule”).

The authors of CheXpert [9] propose an evaluation protocol over 5 categories: “Atelectasis”, “Cardiomegaly”, “Consolidation”, “Edema”, and “Pleural Effusion”, which were selected based on the clinical importance and prevalence from the validation set. In this experiment, we use the official set to train all the models, and show the AUC scores of these 5 abnormalities on official validation set. CheXpert captures uncertainties inherent in radiograph interpretation with an effective labeling strategy ($0$ for negative, $-1$ for uncertain, and $1$ for positive). There are a few significant differences between the performance of the uncertainty approaches. U-Ignore [9] ignores the uncertainty labels during training, while U-Zeros [9] and U-Ones [9] treat them as $0$ or $1$. Since the U-Ones model achieves the best AUC performance in this dataset, we treat all uncertainty labels as $1$ in our experiments. At the same, we show the results of our methods with two backbone networks (ResNet50 and ResNet152). AUC scores are shown in Table 8. Our method outperforms the three models in the official paper [9], indicating that our method can maintain competitive classification performance.

In Table 8, we also compare the performance of DenseNet-121 with the conditional training and label smoothing regularization (U-Ones+CT+LSR) strategies in [36]. Pham et al. [36] obtained mean AUC score of 0.930 on the official CheXpert testing set with the ensemble approach that combined six deep models of DenseNet-121, DenseNet-169, DenseNet-201, Inception-ResNet-v2, Xception, and NASNetLarge. In [36], the DenseNet-121 with the strategies of U-Ones+CT+LSR is the single model that achieves the best performance in the official CheXpert validation set. In comparison, our method can outperform DenseNet-121 with the strategies of U-Ones+CT+LSR even with smaller backbone of ResNet50 on the same validation set. The mean AUC scores on the official test set of our method with the backbone of ResNet50 and ResNet152 are 0.888 and 0.895 respectively without the implementation of an ensemble strategy.

In this experiment, we split 221,674 images from the official training set to train the models, which including 457 images with 1435 bounding boxes. Then, we split 1888 images with 4664 bounding boxes as the validation set to evaluate the localization performance of 9 abnormalities. Table 9 and Table 10 show the localization results in the validation set. Here, we compare with two baseline models, ResNet50 and Dilated ResNet50. Dilated ResNet50 encodes the size of feature map to $\frac{1}{8}$ of the input images and uses LSE pooling. We use CAM [13] to generate the attention maps as the localization results for two baseline methods. We can see from these tables that our method outperforms both these baseline methods. Moreover, with the addition of very few extra annotations (457 images), our method (“Ours_extra”) can gain a significant improvements with respect to localization scores. We also show the results of self-refinement method (“Ours_extra + SR”), which can improve the $T(IoR)$ results.

Qualitative results are shown in Fig. 8, where we observe that our method can generate more accurate localization results. We can see the baseline ResNet50 model usually produces rough and large attention maps, while our method can produce more accurate results. The gridding effect can be observed in Fig. 8 from the results of the Dilated ResNet50 model, which is caused by dilated convolution [37]. We can see that our method can eliminate this effect when using dilated convolution to increase the resolution of attention maps. With the addition of very few extra annotations, we can see that attention results of our method can be further improved. We can also see that the attention results using the self-refinement method are concentrated within the box annotations and look anatomically plausible and appear close to the image segmentation effect.