TLMCM Network for Medical Image Hierarchical Multi-Label Classification

A generic HMC task must respect the hierarchy constraint, i.e. for each label that is predicted to be true, all the ancestor labels as pre-defined in the hierarchy structure must also be predicted to be true Giunchiglia and Lukasiewicz (2020). In this regard, we adopted the two key concepts proposed in Giunchiglia and Lukasiewicz (2020): Max Constraint Module (MCM) and Maximum Constraint Loss (MCLoss).

Formally, for a generic HMC task with a set $S$ of $n$ labels in total, given a label $A$, let $D_{A}$ be the set of all labels which are the descendants of the label $A$ in the hierarchical structure, and a machine learning model $h$ predicts the label $A$ to be true with the probability of $h_{A}$, we then impose the MCM module on top of the output of the model $h$, such that the output of MCM for label $A$ is:

In addition, we can also guide the training process by incorporating the MCM constraint into the loss function of the underlying model $h$. Formally, let $y_{A}$ and $y_{B}$ be the ground truth value of label $A$ and $B$, the loss for label A is

Then the final loss function is defined as

According to Giunchiglia and Lukasiewicz (2020), the novel $\mathrm{MCLoss}$ function could bring benefits to the gradient backpropagation to achieve a lower loss than the standard binary cross-entropy loss function.

For the backbone model, we retain the ResNet50 convolution blocks unchanged and focus on adapting the last fully-connected linear layer. In particular, we replace the original last layer with two fully-connected linear layers with a ReLU activation in between and a Sigmoid after. The hidden dimension is set to 256, and the output dimension is equal to the total number of labels. The Sigmoid layer converts the output scores into probabilities, as required by the MCM module.

The modified ResNet50 output is then fed into the MCM to predict the probability of truth for each label. We use the adapted loss function, $\mathrm{MCLoss}$, as defined in Equation (3). The output are the probabilities for each label to be true. In the specific task of MI-HMC, since the label structure is mono-hierarchical(fully explained in the dataset description section 4.1), we just make our prediction that all labels corresponding to the maximum probability of the MCM output are predicted to be true. The whole architecture design is shown in Figure 1.

There are two common approaches to transfer learning: (1)freezing the convolution module as the fixed feature extractor, and only training the linear classifier(ResNet50 as Fix Feature Extractor, hereafter referred to as RNFFE); (2) fine-tuning the convolution module with the pre-trained weights of the convolution module as well as the linear classifier module(ResNet50 Fine-Tuning, hereafter referred to as RNFT). In the experiment section 4.3, we will explore whether RNFFE or RNFT is more suitable for the MI-HMC task.

The main dataset that we used for this research is the 2009 ImageCLEF edition of the IRMA X-ray dataset Thomas and B. (2009). Each image in the dataset was classified based on the IRMA code Lehmann et al. (2003). A classification code may consist of three or four digits, representing a mono-hierarchical classification structure for the corresponding medical image. A mono-hierarchical label means that the classification hierarchy is a tree structure, thus each child node can only have one parent node.

In this study, our focus was directed toward evaluating the effectiveness of our model on two specific classification codes: anatomical (A) of 110 labels, and directional (D) of 36 labels. The dataset consists of 14410 images in total, and we split it into the train/validation/test set with the ratio of 70:15:15. This dataset choice of the (A) and (D) codes also allowed for a direct comparison with the baseline results^***The ImageCLEF07 datasets from their research are an older version no longer available. However, the ImageCLEF09 datasets share the same chracteristics..

Area Under the average Precision and Recall Curve($AU\overline{(PRC)}$) Giunchiglia and Lukasiewicz (2020); Sorower (2010) is most commonly used in multi-label classification and tasks, and is also the metric that we used to compare with the state-of-the-art baseline results. Additionally, Exact Match Ratio($EMR$) Sorower (2010) is a strict accuracy metric in that the prediction is considered correct only when the set of labels of prediction exactly matches the corresponding set of labels of ground truth. Formally, the $EMR$ is computed as follows:

where, $I$ is the indicator function, $n$ is total number of all labels, $Y_{i}$ is the set of ground true labels for sample $i$, and $Z_{i}$ is the set of predicted labels for sample $i$.

From the application perspective of the MI-HMC task, if a prediction is partially correct, it still deserves some credit. In this respect, we utilize the multi-label classification $HammingAccuracy$ Sorower (2010), which is defined as the following with the same meaning as all notations defined in Eq. 4:

We performed experiments with both RNFFE and RNFT on the MI-HMC tasks of ImageCLEF09A and ImageCLEF09D. We maintained the same learning rate of $5e-6$, weight decay of $1e-6$, batch size of 32 and the total number of epoches of 120 for each task. The ReLU activation was used for the linear classifier and no dropout was applied. Then we trained the model with the Adam optimizer.

The training was completed on the Google Cloud Platform(GCP) virtual machine. We used the regular configuration of the “n1-standard-4” machine type, and 1 NVIDIA T4 GPU. We finished 4 training sessions(2 tasks, 2 approaches) in 10 hours, which is reasonably computationally efficient. The code is at https://github.com/flowing-time/IMAGE-HMC.

The $AU\overline{(PRC)}$, $EMR$, and $hammingAccuracy$ for the tasks of ImageCLEF09A and ImageCLEF09D are summarized in Table 1. On both tasks, the transfer learning approach RNFT shows a significantly higher $AU\overline{(PRC)}$ score than the baseline results of C-HMCNN(h) in Giunchiglia and Lukasiewicz (2020) and some other recent models with good performance(e.g. Wehrmann et al. (2018a), Pelka et al. (2018)).

To the best of our knowledge, the $EMR$ and $hammingAccuracy$ metrics have not been commonly reported on the previous HMC tasks. Nevertheless, our TLMCM network shows 80% - 90% prediction accuracy on the MI-HMC task and demonstrated its great value in real-world healthcare applications.

For a better understanding of fine-turning approaches in our methodology, we plotted the learning curves of both approaches for the ImageCLEF09A task in Figure 2. The logarithmic scale is used for the loss(1st colummn), while linear scale is used for the metrics(2nd, 3rd, 4th column). It is clearly shown that significantly lower training and validation loss can be achieved by the RNFT approach than RNFFE. In the 2nd column, we can see RNFT quickly attains the plateau of the highest $AU\overline{(PRC)}$ score, for both the training and validation set. The 3rd and 4th columns show the prediction accuracy($EMR$ and $HammingAccuracy$) over the training iterations and the RNFT approach is again the winner.

Overall, the experiments show that for the MI-HMC task, our TLMCM network is highly effective, and the transfer learning approach of fine-tuning with pretrained weights(RNFT) should be the prioritized choice, as it does not only perform best in all three metrics but also achieves the best result with the fewest training iterations.