ProCo: Prototype-aware Contrastive Learning for Long-tailed Medical Image Classification

Classic contrastive training pairs (i.e., positive and negative pairs) are used to learn the representation of instances. However, in the long-tailed dataset, the head classes dominate most of negative pairs via the conventional contrastive methods, causing the under-learning of tailed classes. Previous works [2, 24] reveal that not all negative pairs facilitate the contrastive learning. Therefore, the key to improve the performance of the long-tailed problem is to reduce the redundancy of contrastive negative pairs and mine recognizable positive pairs. To this end, we propose a new concept named category prototype. The prototype is a set of learnable parameters $P=\{p_c\in R^k|1\le c\le C\}$ for predefined $C$ categories and is optimized by our proto-loss as described in Sec. 3.3. Then, we generate the adversarial proto-instance from the category prototype and representative sample, i.e., confusing samples from the alternative sample queue, via a linear interpolation in the feature space, which will be utilized to form a training pair in ProCo. Theoretically, the adversarial proto-instance is designed as a special outlier, which can encourage ProCo to rectify the decision boundaries of the tailed categories during the contrastive learning.

Fig. 2 illustrates the diagram of the adversarial proto-instance. First, for each instance $(x_j,y_j)$, the queue $Q$ is divided into two disjoint subsets $Q^{+}=\{s_i|y_i=y_j\}$ and $Q^{-}=\{s_i|y_i\ne y_j\}$ that comprise positive and negative instances, respectively. Similarly, $P$ is also grouped into a singleton set $P^{+}=\{p_c^{+}|c=y_j\}$ with exactly one prototype $p_c^{+}$ and a negative set $P^{-}=P\setminus P^{+}$, where $p_c^{+}$ is the positive prototype feature corresponding to the current instance. Then, the adversarial positive proto-instances are derived from $P^{-}$ and $Q^{+}$ while the negative ones are from $P^{+}$ and $Q^{-}$.

To synthesize the adversarial negative proto-instances, we prioritize those negative instances that are likely to be confused with the current instance. The distance between the current feature $G\left(x_j\right)$ and negative feature $s_i^{-}\in Q^{-}$ can be utilized as an indicator, and is represented as:

where the superscript $\top$ represents the transpose operation. Then, we rank the negative instances in $Q^{-}$ in ascending order according to their distance to $G\left(x_j\right)$ and select the top $\gamma$ instances to compose the adversarial proto-instance set $Q_{\gamma }^{-}$:

where $s_{\gamma }$ is the $\gamma$-th element in the sorted $Q^{-}$. Finally, for more challenging negative instances, we randomly perturb each element in $Q^{-}$ with the positive prototype $p_c^{+}$:

where ${ϵ}_i\in (0,E)$ is a random interpolation coefficient for each sample and the upper bound $E$ is a hyperparameter with a small value. We assume that prototype $p_c^{+}$ always contributes less than $s_i$ to the generated proto-instance, which guarantees that the negative semantic within $s_i$ can be held.

As for the adversarial positive proto-instances, the strategy is to select samples in $Q^{+}$ that are misclassified and combine them with the incorrectly assigned prototype, where the interpolation manner is identical to that of Eq. 3.

Although the problem of constructing contrastive pairs has been addressed by the proposed category prototype and adversarial proto-instance, the learning of category prototype is still potentially affected by the class imbalance problem. We argue that the underlying reason is the prototype bias, where the generated proto-instance will incline to head classes, jeopardizing the performance on tailed classes. For this reason, this work proposes a prototype recalibration strategy, which estimates the representative level of each category prototype showing the importance of tailed class prototype features.

Specifically, inspired by the similarity between the projected features and prototype, we introduce a rectified sigmoid function to achieve recalibration. In formal, the calibration factor for each category prototype $p_c\in P$ is defined as:

where $\{x_j|y_j=c\}$ is a subset of samples in the associated category $c$, and $N_c$ is the total number of them. \colorblue To allow end-to-end calibration for batch-based training, we keep a running mean via the exponential moving average for the global calibration factor:

where $\{x_j|y_j=c\}$ are samples in a batch with label $c$, and $\beta$ is a smoothing coefficient.

The calibration factor reflects the difficulty of each category and the representativeness of the corresponding prototype in the model learning. Eventually, we impose all prototypes in $P$ by the calibration factors as:

To integrate the contrastive learning into our work in an end-to-end manner, we refer to the concept of InfoNCE[22]. Unfortunately, it only supports single positive pair, which is incompatible to our work. Therefore, inspired by the unified contrastive loss [6], we extend it to include prototypes and involve both positive and negative adversarial proto-instance for training, which ensures the optimization consistency of supervised and contrastive training so as to achieve decent performance compared with the former studies.

Formally, considering $H^{-}=Q^{-}\cup {\hat{Q}}_{\gamma }^{-}$ as negative set, and $H^{+}=Q^{+}\cup {\hat{Q}}_{\gamma }^{+}$ as positive set. For an instance $x_j$, the proto-loss is formulated as:

Note that the proposed proto-loss can also be applied to general contrastive settings.

We conduct the experiment on two publicly available datasets. The ISIC2018 is accessed by the Skin Lesion Analysis Toward Melanoma Detection 2018 challenge [5]. The other dataset is APTOS2019, which is provided by  [1]. For a better illustration, the details of datasets are listed in Table 1. Notably, the imbalance ratio denotes as $N_{max}/N_{min}$, where $N$ is the number of samples in each class. We follow the same protocol in [21] and randomly split the original dataset into train and test sets with the ratio of 7:3. All experimental results will be reported with the criterion of accuracy and F1-score.

In this work, we selected 9 existing methods as the comparison methods. To be specific, we use cross-entropy (CE) as the baseline. Further, based on CE, the balanced resampling strategy is used to address the imbalanced classification problem. The contrastive learning method can be treated as another comparison method. By integrating the balanced resampling strategy,“CL+resample” is proposed in [21], which is a two-stage approach by training the backbone firstly, then freezing the backbone and training the classifier with balanced sampling. To down-weight easy negatives in one-stage detector, focal loss [18] is also useful in classification problem. LDAM [3] focuses on label distribution margin, and is regarded as a simple but effective training strategy. OHEM [26] is a hard negative mining method based on the model. DANIL [9] explores distractors to learn better CNN features.

The data augmentation policy and update ratio we utilized of the contrastive learning is identical to MoCoV2 [4]. ResNet50 [11] is used as our backbone. We implement projector $G$ with 2 fully-connected layers, of which the hidden layer size is set to 2048, followed by ReLU activation function. The projector $F$ is realized by a single fully-connected layer with ReLU activation function. For simplicity, category prototype $P$ can be regarded as a classifier in our framework. The batch size is 128 and the default optimizer is SGD with a momentum of 0.9 and a weight decay of 0.0001. The initial learning rate is set to 0.05. The initial similarity value ${\omega }_c$ is set to 0.01 and $\beta$ is set to 0.95. Hyperparameters $\gamma$ and $E$ are set at 20 and 0.4 via a grid search. Referring to [21], the training epochs of ISIC2018 and APTOS2019 are set at 1,000 and 2,000, respectively. Particularly in the test phase, we only utilize $F\left(x\right)$ and $P$ to acquire the prediction.

In this part, we compare the proposed ProCo with the state-of-the-art methods on two open-release datasets: ISIC2018 and APTOS2019. Table 1 presents the entire experimental results. We can see that the proposed method achieves optimal performance regardless of the data set, demonstrating its excellent generalization. Apart from this observation, the performance of Focal loss and LDAM is comparable to the baseline on both datasets. In addition, OHEM, MTL, and DANIL obtain comparable accuracy and F1-score on the APTOS2019 while inferior performance on the ISIC2018, illustrating their weak generalization. The results from the two CL-based methods are uniform and outperform the baseline. Note that the role of the balanced resampling strategy is inconsistent across different data sets.

In this work, the proposed framework has three fundamental modules. To validate the effectiveness of each module, we carried out ablation studies as shown in Table 3. For this ablation study, four extra experiments have been designed by arranging and combining these three modules: 1) We discarded the proto-instance and prototype recalibration strategy and only used Proto-loss to train the proposed method. 2) Based on the above experiment setting, we introduced the proto-instance module to re-train the proposed method. 3) Integrating the prototype recalibration strategy into the first experiment setting. 4) Employing all the three modules which is our entire proposed framework. From Table 3, we can clearly observe that with all the three modules, the proposed method can obtain the best performance in terms of accuracy and F1-score as shown in the last row. To be specific, the modified framework with only the Proto-loss as the loss function obtained an inferior performance to other three experiments. To this end, we can benefit from proto-instance module and prototype recalibration strategy, which is consistent with the results from experiment 1 and 2. In addition, to evaluate the superiority of the Proto-loss over other commonly-used loss functions: cross-entropy (CE) and InfoNCE, we re-trained the proposed network using different loss functions, respectively. The final results of F1-score were: 0.716 (CE), 0.739 (InfoNCE), and 0.742 (Proto-loss). We can see that using Proto-loss as the loss function will improve the final F1-score by 3.63% compared to that from CE.