Comparing to Learn: Surpassing ImageNet Pretraining on Radiographs By Comparing Image Representations

ImageNet [2] pretraining has been proved to be an effective way to perform 2D transfer learning for medical image analysis. Lots of experiments have shown that, compared with learning from scratch, pretrained models not only help to achieve higher accuracy but also speed up the model convergence. These benefits can be attributed to two factors: (a) effective learning algorithms designed for deep neural networks and (b) generalized feature representations learned from a great quantity of natural images. However, there exists an obvious domain gap between natural images and medical images, which raises a question whether we can have a pretrained model directly from medical images and how to approach this.

As is well known, we need domain experts’ diagnosis to produce reliable medical annotations, which definitely help improve our model performance. On the other side, it is often difficult to access to a great number of doctors’ conclusions considering limited medical resources as well as protection of patient privacy. So how to develop algorithms to learn from a vast amount of data without annotations has drawn more attention in the medical imaging community. Zhou et al. [13] proposed a self-supervised pretraining method Model Genesis which utilized medical images without manual labeling. On the chest X-ray classification task, Model Genesis is able to achieve comparable performance with ImageNet pretraining but still cannot beat it.

In this paper, we present a novel self-supervised pretraining approach focusing on providing pretrained 2D deep models for radiograph related tasks from massive unannotated data. We name our method as Comparing to Learn (C2L) because the goal is to learn general image representations by comparing different image features as the supervision. Different from Model Genesis [13] that resorts to an image restoration pretext task, the supervision signal of the proposed C2L comes from the self-defined representation similarity. Similar ideas have been adopted in [1, 9, 4], where most of them take advantage of the transitive invariance of images to produce self-supervised signal. On the contrary, in this paper, we mainly focus on feature level contrast and propose to construct homogeneous and heterogeneous data pairs by mixing image and feature batches. Moreover, a momentum-based teacher-student architecture is proposed for the contrastive learning, where the teacher and student networks share the same structure but are updated differently. To be specific, the teacher model is updated using both itself and the student network. Extensive experiments of different datasets and tasks demonstrate that the proposed C2L method can surpass ImageNet pretraining and other competitive baselines by non-trivial margins.

In this section, we introduce the proposed Comparing to Learn (C2L) method in details. The overall workflow is provided in Figure 1.

Batch mixup and feature mixup. As shown in Figure 1 and Algorithm 1, for each input image batch, we first use random augmentation (e.g., random cropping, rotation, and cutout [3]) to generate two augmented batches. Different from traditional image-level mixup [11], a batch-wise mixup operation is proposed to apply to each augmented batch. Suppose each batch $\mathcal{X}^{A}_{i}$ contains $Z$ images where $i=\left\{1,2\right\}$, we randomly shuffle $\mathcal{X}^{A}_{i}$ to construct its paired batch $\mathcal{X}^{M}_{i}$ which can be expressed as:

where $\lambda\sim\text{Beta}(1.0,1.0)$ and Beta stands for the beta distribution. In practice, we found that using the same mixing factor $\lambda$ and shuffling method for both batches ($i=\left\{1,2\right\}$) actually helps improve the model performance. As for feature mixup in Figure 1, we apply the same mixing strategy to the feature representations.

Teacher network. An intuitive idea is to use the same model for both student and teacher networks. However, we found that such strategy does not work in practice which may lead to gradient explosion. Meanwhile, constructing teacher network using momentum update has been widely adopted as a way to produce stable predictions [8, 4]. In our case, using momentum helps stabilize the training process and reduce the difficulty of network optimization. As shown in line 17 from Algorithm 1, the momentum function can be formalized as:

where we use an exponential factor $\theta$ to control the degree of momentum. We can see that the teacher model $\text{net}_{t}$ is updated using both itself and the student network $\text{net}_{s}$. In practice, we pass $\mathcal{X}_{1}^{A}$ and $\mathcal{X}_{1}^{M}$ to the student network while $\mathcal{X}_{2}^{A}$ and $\mathcal{X}_{2}^{M}$ are passed to the teacher network. In Algorithm 1, we use $v_{1}^{\left\{A,M\right\}}$ to represent feature vectors from the student model and $v_{2}^{\left\{A,M\right\}}$ are used to represent the outputs of the teacher model.

Homogeneous and heterogeneous pairs. For constructing homogeneous pairs, we assume that applying data augmentation (including mixup operation) only slightly change the distribution of training data. Based on this, each homogeneous pair should contain the results after the same set of operations which includes random augmentation, batch mixup or feature mixup as shown in Figure 1. For heterogeneous pairs, we simply contrast the current features with all preceding features stored in the memory queue.

Feature comparison, memory $Q$ and loss function. As we have mentioned above, the goal of C2L is to minimize the distance between homogeneous representation pairs such as $(v^{A}_{1},v^{A}_{2})$ and $(v^{M}_{1},v^{M}_{2})$. Meanwhile, it is also necessary to maximize the difference between heterogeneous representations, in which we contrast current features with past features which are collected from past training iterations. To store these past representations, we employ a memory queue $Q$ proposed in [4]. The reason why we use a large $Q$ is that we hope to contrast current features with a great number of preceding features because more comparisons usually lead to better representations (as shown in Table 2). So the pairs of current features and past features can be formalized as ($v^{A}_{1}$, $Q$) and ($v^{M}_{1}$, $Q$), as shown in Figure 1. For simplicity, we use $v^{Q}_{i}$ to denote a specific feature vector in $Q$ where $i=\left\{1,2,...,N\right\}$ and $N$ is the length of queue. To be specific, we import a subscript $j$ to index a feature vector given a batch of features. Then, we can convert this distance measurement problem to a naive classification problem. For image $j$ in $\mathcal{X}^{A}_{1}$, $\{(v_{1,j}^{A})^{T}v_{2,j}^{A},(v_{1,j}^{A})^{T}Q\}$ in line 12 can be further expressed as:

whose length is $N+1$. A similar case also exists in $\mathcal{X}^{M}_{1}$. For more comparison, we apply feature mixup to $v^{A}_{2}$. In Algorithm 1, we use $v_{m}$ to represent the output of feature mixup. Similarly, we also compare $v^{M}_{1}$ with $v_{m}$ and $Q$ which leads to another set of $N+1$ predictions:

It is worth noting that the first item in each set should be larger than other items because the first item is an inner product of homogeneous representation. Thus, we can apply a cross entropy loss (CE) to the above sets of predictions where a one hot vector {1,0,0,…,0} is treated as the ground truth for each set.

After we update the network parameters, we also update $Q$ by inserting $v_{2}^{A}$, $v_{2}^{M}$, and $v_{m}$, respectively. Since $Q$ is a queue and has a fixed size, the previous feature vectors are automatically removed. After we complete the training stage, only $\text{net}_{s}$ is extracted to become a pretrained model.

For pretraining, we employ C2L to pretrain ResNet-18 and DenseNet-121. The default batch size is 256 and the size of each input image is 224$\times$224. For input augmentation, we apply random crop, rotation (10 degree), grayscale and horizontal flip to each input batch. Moreover, we also add cutout to augmented images in order to increase the diversity of transformation for learning better representations. We use L2 normalization for each feature vector. The momentum factor $\theta$ is set to 0.999 and the length of queue $Q$ is $2^{15}=32768$. We use SGD as the default optimizer where the initial learning rate is 0.03 and its weight decay is 0.0001. We train each model for 240 epochs and the learning rate is divided by 10 at 120, 160 and 200 epochs, respectively. When fine-tuning pretrained models on ChestX-ray14 and CheXpert, the input image size is set to 224$\times$224 and we train both ResNet-18 and DenseNet-18 for 50 epochs. As for Kaggle Pneumonia Detection, we employ RetinaNet which uses ResNet-18 as backbone. The default image size is 512$\times$512 and the batch size is 4.

In this part, we conduct experiments on ChestX-ray14. As we have mentioned above, we use 70k unlabeled images for C2L pretraining and then use the rest labeled training set for fine-tuning. We report averaged AUROC performance and results of eight class are also provided (results of all fourteen classes are provided in the supplementary material). Note that to save space, the ablation study of $\theta$ (cf. Equation 2) is put in the supplementary material.

We first report the ablation results of proposed mixup approaches in Table 1. We can see that the proposed batch mixup can already outperform the original mixup method [12] by 0.6 point on average performance. This is because using the same $\lambda$ for shared batches may help maintain the consistency between batches. After adding mixed consistency loss, we can improve the batch mixup method by about 0.3 point. Since the goal of C2L is to learn powerful feature representations, we further apply mixup to generated features. Somewhat surprisingly, we can find that the propose feature mixup can be well integrated with $\text{loss}_{M}$ and surpass batch mixup by approximate 1 point. In summary, the proposed mixup strategies can outperform the original mixup method by 1.6 points.

Another important component of C2L is the augmentation strategies. An appropriate augmentation method should reasonably increase the diversity of augmented batches. Such characteristic may help pretrained models to learn representations which are discriminative enough to distinguish different radiographs. In Table 2, we investigate the effects of widely adopted augmentation strategies. It is normal to find that random rotation and cutout can enhance the performance by 0.3 point while adding jigsaw and dropout may degrade performance. Moreover, we find that simply increasing the length of $Q$ can be harmful. We argue the reason is that a longer queue may contain more useless features and thus reduce the attention of other useful representations.

In this section, we compare C2L pretrained models with Model Genesis, ImageNet pretraining and MoCo [4]. For pretraining datasets, we merge the training sets of ChestX-ray14 and CheXpert with radiographs in MIMIC-CXR and MURA to generate an unlabeled database containing approximate 700k images. For network architectures, we deploy ResNet-18 and DenseNet-121, both of which are widely used networks.

ChestX-ray14. We report fine-tuned AUROC results on the validation set. Besides ImageNet pretraining, we also perform experiments using Model Genesis (MG) [13] and recently proposed MoCo [4]. In Table 3, the proposed C2L method surpasses other approaches by a significant margin. In fact, although MG and MoCo are able to achieve comparable results with ImageNet pretraining, they cannot surpass ImageNet pretrained model significantly. However, C2L pretrained model outperforms ImageNet pretraining by 2 points on ResNet-18. On DenseNet-121, C2L achieves 84.4% averaged AUROC which is 1.5 points higher than ImageNet pretraining.

CheXpert. Similar to ChestX-ray14, we fine-tune the pretrained models using the training set. We can see that ImageNet pretraining performs better than MG while MoCo achieves comparable results. In contrast, C2L generates better pretrained models. On ResNet-18, C2L outperforms ImageNet pretraining by 1.2 points. When it comes to DenseNet-121, the performance gap becomes 1.4 point.

Kaggle Pneumonia Detection. We use ResNet-18 as the backbone of RetinaNet. In Table 5, it is obvious that C2L outperforms ImageNet pretraining on all thresholds significantly, especially on large thresholds. As for MoCo and MG, MoCo is marginally better than ImageNet pretraining while MG performs slightly worse.

We proposed a self-supervised pretraining method C2L (Comparing to Learn) to learn medical representations from unlabeled data. Our approach makes use of the relation between images as supervision signal and thus requires no extra manual labeling.

This work was funded by the Key Area Research and Development Program of Guangdong Province, China (No. 2018B010111001), National Key Research and Development Project (2018YFC2000702) and Science and Technology Program of Shenzhen, China (No. ZDSYS201802021814180).