Jointly Boosting Saliency Prediction and Disease Classification on Chest X-ray Images with Multi-task UNet

Abstract

Human visual attention has recently shown its distinct capability in boosting machine learning models. However, studies that aim to facilitate medical tasks with human visual attention are still scarce. To support the use of visual attention, this paper describes a novel deep learning model for visual saliency prediction on chest X-ray (CXR) images. To cope with data deficiency, we exploit the multi-task learning method and tackle disease classification on CXR simultaneously. For a more robust training process, we propose a further optimized multi-task learning scheme to better handle model overfitting. Experiments show our proposed deep learning model with our new learning scheme can outperform existing methods dedicated either for saliency prediction or image classification. The code used in this paper is available at [webpage, concealed for double-blind review].

Appendices

AMathematical Derivation of Vicious Circle for Overfitting

Let $L\ge 0$ be the loss for a task, $\mathcal {T}$, and $\sigma >0$ be the variance estimator for L used in Eq. 1. Therefore, the loss for $\mathcal {T}$ following Eq. 1 can be expressed as:

$$\begin{aligned} \mathcal {L} = \frac{L}{\sigma ^2}+\ln (\sigma +1). \end{aligned}$$

(6)

The partial derivative of $\mathcal {L}$ with respect to $\sigma$ is:

$$\begin{aligned} \frac{\partial \mathcal {L}}{\partial \sigma } = -\frac{2L}{\sigma ^3}+\frac{1}{\sigma +1}. \end{aligned}$$

(7)

During a gradient based optimization process, to minimize $\mathcal {L}$, $\sigma$ converges to the equilibrium value ($\sigma$ remains unchanged after gradient descend) which is achieved when $\frac{\partial \mathcal {L}}{\partial \sigma }=0$. Therefore, the following equation holds when $\sigma$ is at its equilibrium value, denoted as $\tilde{\sigma }$:

$$\begin{aligned} L = \frac{\tilde{\sigma }^3}{2\tilde{\sigma }+2} \end{aligned}$$

(8)

which is calculated by letting $\frac{\partial \mathcal {L}}{\partial \sigma }=0$. Let $f(\tilde{\sigma }) = L$, $\tilde{\sigma }>0$, we can calculate that:

$$\begin{aligned} \frac{d f(\tilde{\sigma })}{d \tilde{\sigma }} = \frac{\tilde{\sigma }^2(2\tilde{\sigma } + 3)}{2(\tilde{\sigma } +1)^2}>0, \quad \forall \tilde{\sigma }>0. \end{aligned}$$

(9)

Therefore, we know that $f(\tilde{\sigma })$ is strictly monotonically increasing with respect to $\tilde{\sigma }$, and hence the inverse function of $f(\tilde{\sigma })$, $f^{-1}(\cdot )$, exists. More specifically, we have:

$$\begin{aligned} \tilde{\sigma } = f^{-1}(L). \end{aligned}$$

(10)

As a pair of inverse functions share the same monotonicity, we know that $\tilde{\sigma } = f^{-1}(L)$ is also strictly monotonically increasing. Thus, when L decreases due to overfitting, we know that $\tilde{\sigma }$ will decrease accordingly, forcing $\sigma$ to decrease. The decreased $\sigma$ leads to an increase in the effective learning rate for $\mathcal {T}$, forming a vicious circle of overfitting.

B Training Settings

We use the Adam optimizer with default parameters [29] and the RLRP scheduler for all the training processes. The RLRP scheduler reduces $90\%$ of the learning rate when validation loss stops improving for P consecutive epochs, and reset model parameters to an earlier epoch when the network achieves the best validation loss. All training and testing are performed with the PyTorch framework [45]. Hyper-parameters for optimizations are learning rate r, and P in RLRP scheduler. The dataset is randomly partitioned into $70\%$, $10\%$ and $20\%$ subsections for training, validation and testing, respectively. The random data partitioning process preserves the balanced dataset characteristic, and all classes have equal share in all sub-datasets. All the results presented in this paper are based on at least 5 independent training with the same hyper-parameters. NVIDIA V100 and A100 GPUs (Santa Clara, USA) were used.

C Saliency Map visualization

Table 4. Visualization of predicted saliency distributions. The ground truth and predicted saliency distributions are overlaid over CXR images. Jet colormap is used for saliency distributions where warmer (red and yellow) colors indicate higher concentration of saliency and colder (green and blue) colors indicate lower concentration of saliency.