Dual-channel Prototype Network for few-shot Classification of Pathological Images

The global feature modeling capability of Vision Transformers has been demonstrated to hold significant advantages in multiple sub-domains of computer vision. Among these, the Pyramid Vision Transformer (PVT) [28], as an improved version of ViT, aims to reduce the computational complexity of multi-head self-attention by introducing a pyramid structure, thereby capturing richer contextual features. Hence, in this paper, the PVT network is chosen as the encoder for global feature representation of pathological images.

The Masked Autoencoder (MAE) [29] exhibits outstanding performance across multiple visual tasks through its asymmetric encoder-decoder architecture. This achievement is partly attributed to the Vanilla Vision Transformer’s ability to extract global features. However, the PVT network typically introduces operations within local windows to reduce the quadratic complexity of global self-attention, with the possibility of complete loss of visual token information within these randomly constituted visual token sequences. To address this issue, our study introduces a uniform masking strategy [30], which improves upon the random masking method in MAE, enabling the effective application of the MAE architecture to the pre-training of PVT.

Specifically, Uniform masking includes two stages of masking: Uniform Sampling (US) and Secondary Masking (SM). As illustrated in Figure 3, the original input image is first partitioned into chunks; then, during the US phase, 25% of the patches are sampled from the image based on a uniform distribution strategy, ensuring that at least one patch is selected in every 2×2 grid across the entire image, resulting in a uniformly distributed set of sparse image blocks. Following this, SM further randomly masks a portion (about 25%) of the sampled areas without completely discarding the masked blocks. Instead, these are treated as a shared, learnable token sequence and combined with visible visual tokens as input to the encoder.

The design of the encoder, decoder, and reconstruction target follows the original setup of MAE. However, unlike the original MAE, the encoder here is the PVT, which contains four stages, each outputting feature maps of different resolutions with a total stride of 32. This pyramid-structured transformer encodes the input 16*16 patches into sub-pixel level latent representations. Therefore, before inputting into the decoder, a PixelShuffle operation is required to restore the resolution of the latent features. Finally, the decoder reconstructs the missing pixels of the original image and uses the Mean Squared Error (MSE) as the loss function to optimize the model, with the MSE formula as follows:

where $N$ is the total number of missing pixels in the original image, $Y_{i}$ is the true pixel value of the $i$-th pixel, and $\hat{Y}_{i}$ is the predicted pixel value by the model. By minimizing the MSE loss, the model is trained to accurately reconstruct each pixel, thereby learning the global features of the image.

The multi-scale prototype representation is achieved through a dual-channel encoder, as shown in Figure 4(a). We utilize a pretrained PVT to extract the global features of pathological images, while a CNN is employed to extract local features. Specifically, for a given input image $x$:

Here, $Z_{G}$ and $Z_{L}$ represent the global and local feature embeddings of the input image $x$, respectively, each with a dimension of $D$. To further refine these features and reduce their dimensionality, we apply Principal Component Analysis (PCA) [31] to process the global and local feature embeddings, extracting a set of linearly independent feature embeddings and removing redundant information as much as possible. The processed features are then concatenated, resulting in a mixed semantic feature:

Here, $Z^{\prime}_{G}$ and $Z^{\prime}_{L}$ are the global and local features post-PCA dimensionality reduction, each with a dimension of $D/2$. $Z_{\text{Mix}}$ is the mixed semantic feature obtained by concatenating $Z^{\prime}_{G}$ and $Z^{\prime}_{L}$, with a dimension of $D$. At this point, the input image $x$ can be represented as a collection of three different scale feature sets $\{Z_{G},Z_{L},Z_{\text{Mix}}\}$.

Next, to construct the multi-scale prototype representation, we consider a support set $D_{\text{support}}$ containing $N$ categories, with each category $C_{N}$ comprising a set of images $\{x_{1},x_{2},...,x_{K}\}_{i=1,2,..,K}$. Here, $K$ denotes the number of samples per category in the support set, i.e., K-shot. Each image is mapped to the multi-scale feature space, resulting in $\{Z^{i}_{G},Z^{i}_{L},Z^{i}_{\text{Mix}}\}_{i=1,2,..,K}$. For each category $C_{N}$, we compute its multi-scale prototype representation in the respective feature spaces, which is the mean of all image features in the support set:

Here, MP denotes the multi-scale prototype. Thus, each $C_{N}$ category is represented by a set of multi-scale prototype representations $MP_{C_{N}}=\{P_{G}^{C_{N}},P_{L}^{C_{N}},P_{\text{Mix}}^{C_{N}}\}$. Ultimately, we obtain a multi-scale prototype representation matrix:

This matrix reflects the average features of each category sample in the support set across different scales, aiming to provide a rich and effective prototype representation for the task of pathological image classification in FSL scenarios.

The soft voting strategy, a common classification method in ensemble learning [26, 32, 33], typically relies on voting based on the probability predictions of multiple classifiers. This approach fully considers the confidence of each classifier and can enhance the performance of the classifier. Unlike traditional soft voting methods, in this paper, we adopt a soft voting strategy from the perspective of multi-scale features on a single classifier to avoid the overfitting problem caused by training multiple classifiers, as illustrated in Figure 4(b).

Given a query set sample $x_{q}$, we first extract a set of multi-scale feature sets $\{Z_{G}^{q},Z_{L}^{q},Z_{Mix}^{q}\}$ through a dual-channel encoder. To determine the category of $x_{q}$, we calculate the Euclidean distance between the query sample and the prototypes of each category in the multi-scale prototype representation matrix at the corresponding scale, as follows:

where $d_{C_{N}}^{G}$, $d_{C_{N}}^{L}$, and $d_{C_{N}}^{Mix}$ respectively represent the Euclidean distances between the global, local, and mixed features of the query sample $x_{q}$ and the prototype representation of the $C_{N}$th category at the corresponding scale.

To convert the distances into confidences, we use a negative exponential function to prevent numerical instability caused by distances approaching zero. The confidence level for each scale is calculated through the negative exponential of the distance:

Subsequently, we normalize these confidences using the softmax function to calculate the probability distribution of the query sample $x_{q}$ belonging to each category:

Finally, we select the category with the highest probability as the prediction result:

where $y$ represents the predicted category. The optimization process for model training employs the following loss function:

where $y_{q}$ is the true label corresponding to $x_{q}$.

Our study selected three public histological datasets for model evaluation, with details as follows:

• •

  CRCTP dataset [34]: This dataset includes 280,000 non-overlapping patches of size 150$\times$150, extracted at a 20$\times$ magnification from 20 H&E stained colorectal cancer Whole Slide Images (WSIs), encompassing seven different tissue types: Tumor, Stroma, Complex Stroma, Muscle, Debris, Inflammatory, and Benign. Within this dataset, 70% is used as the training set, and 30% as the test set.

• •

  NCTCRC dataset [35]: Comprising 107,180 non-overlapping patches of size 224$\times$224, extracted from 86 H&E stained colorectal cancer WSIs, this dataset includes nine different tissue types: Adipose (ADI), background (BACK), debris (DEB), lymphocytes (LYM), mucus (MUC), smooth muscle (MUS), normal colon mucosa (NORM), cancer-associated stroma (STR), and colorectal adenocarcinoma epithelium (TUM). Here, 100,000 patches are designated for training, with the remaining 7,180 for testing. Although both NCTCRC and CRCTP originate from colorectal sources, they differ in disease categories and regional data sources.

• •

  LC25000 dataset [36]: This dataset contains 25,000 patches of size 768$\times$768, with five different tissue types: colon adenocarcinomas, benign colonic tissues, lung adenocarcinomas, lung squamous cell carcinomas, and benign lung tissues, with 5,000 patches for each tissue type. In this dataset, 70% is allocated for training and 30% for testing.

The diversity in disease and source among these three datasets allows us to construct tasks with varying degrees of domain shift. Specifically, we define three different data realms based on the origin of new classes and base classes: when new and base classes come from the same organ and data source, it is considered same domain. When new and base classes originate from the same organ but different data sources, it is categorized as near domain. When new and base classes derive from not entirely the same organ and different data sources, it is defined as mixture domain. The specific task settings are as follows:

1. 1.

   Same-domain task: New and base classes both come from the CRCTP dataset, defining a 7-way n-shot task. In this task, meta-training and meta-testing are carried out on the training and testing sets of the CRCTP dataset, respectively.

2. 2.

   Near-domain task: Utilizing CRCTP as the base dataset for meta-training, and NCTCRC as the new class dataset for meta-testing, specifically a 5-way n-shot task.

3. 3.

   Mixture-domain task: Similar to the Near-domain task, CRCTP is used as the base dataset for meta-training, but LC25000 is chosen as the new class dataset for meta-testing, specifically a 5-way n-shot task. It is noteworthy that within LC25000, two classes are related to colorectal tissues, and three to pulmonary tissues, thereby defining it as a mixture-domain.

We employed the t-SNE method to reduce the datasets of near domain and mixture domain tasks to three-dimensional space for visualization, observing the data distribution under different tasks, as shown in Figure 5. To quantitatively assess the difficulty of tasks, we calculated the Euclidean distance between datasets. Specifically, the Euclidean distance between CRCTP and NCTCRC datasets in the near domain task is 0.20; while in the mixture domain task, the distance between CRCTP and LC25000 datasets is 1.01. This indicates that the task difficulty of the mixture domain is significantly higher than that of the near domain, consistent with our task design expectations.

During the pre-training phase, our study utilized the UM-MAE algorithm to pre-train a PVT as the encoder for global features on the training sets of the aforementioned three datasets, aiming to mitigate the difficulty of training on FSL tasks. The size of the input images was adjusted to 224$\times$224, with a batch size set to 256 and an efficient batch size of 1024. Optimization was carried out using the AdamW optimizer with betas=(0.9, 0.95), a learning rate of $1\times 10^{-3}$, and epochs set to 100. The remaining parameters were kept consistent with the original UM-MAE paper.

In the FSL phase, the size of the input images remained at 224$\times$224, with a batch size of $N\times K$ (depending on the specific $N$-Way $K$-shot task), and optimization was performed using the Adam optimizer with a learning rate of $1\times 10^{-3}$ and epochs set to 100. For the evaluation phase, 1000 meta-tasks were randomly drawn for each task to assess the model, with each meta-task utilizing 15 samples as a query set and presenting accuracy as the evaluation criterion. All results are the mean of the outcomes from 1000 meta-tasks. To address the imbalance in the number of base and new classes, we followed the settings detailed in [37]. The evaluation metrics involved in this paper are as follows:

where TP, TN, FP, FN respectively stand for True Positives, True Negatives, False Positives, False Negatives.

Our experiment compared various SOTA FSL methods to validate the performance of the proposed approach. This included four classical FSL algorithms: Matching Networks [13], ProtoNet [14], MetaOpt [38], and MAML [8]; as well as three pathology-specific FSL algorithms: Latent Augmentation [39], Histology Siamese Network [19], and Multi-ProtoNets [21]. The experimental results are presented in Table I.

Firstly, in the task of few-shot classification of histological images, DCPN significantly outperformed other classical FSL methods in accuracy. For instance, in the same-domain 10-shot scenario, DCPN achieved an accuracy of 84.67%, which is at least a 6.28% improvement over Matching Networks at 74.69%, ProtoNet at 78.39%, MetaOpt at 74.21%, and MAML at 71.32%.

Compared to the three pathology-specific methods, DCPN still attained the highest accuracy or was on par in all three scenarios. For example, DCPN showed a 2.26% increase in accuracy over Latent Augmentation in the same-domain 10-shot scenario. The Latent Augmentation method primarily increases data diversity by introducing latent variables during training, allowing the model to better adapt to new, unseen classes. However, the randomly generated latent variables may introduce noise that can diminish the representation capability of feature embeddings, limiting further performance enhancements in FSL tasks. DCPN’s advantage was even more pronounced over the Histology Siamese Network, with an accuracy improvement of 5.24%. The Histology Siamese Network applies a deep Siamese neural network directly to the task of few-shot classification of pathological images, learning inherent characteristics that map to class distance on the source domain dataset. However, due to the lack of in-depth modifications tailored to the characteristics of pathological images, its actual effectiveness is limited. Multi-ProtoNets, constructing multiple prototype representations through k-means, achieved the best performance in the same-domain and mixture-domain 10-shot scenarios, but since it does not consider features at different scales, DCPN still achieved the best performance in most scenarios.

The DCPN algorithm, which utilizes multi-scale features, generally surpassed other methods and could even match the performance of supervised training with a full dataset (as shown in Table III). This indicates that using a dual-channel network for feature mining can improve prototype representations, and also proves the importance of multi-scale features for pathological image classification tasks. Furthermore, we observed a significant decrease in network classification accuracy with the increase in domain difficulty. This decline in model generalization capability is due to the growing disparity in data distribution between different domains, as demonstrated in Section 4.1.

The experiment aimed to compare the classification performance of different backbones on a pathological image dataset to determine the optimal feature extraction module, thereby supporting the design of a dual-channel feature embedding network. The CRCTP dataset (base dataset) was used for the training and evaluation of the models. In this experiment, we assessed 11 different backbones, including 6 classic CNN models and 5 Transformer models. The overall experimental results are summarized in Table II. The findings indicate that within the CNN category, ResNet50 achieved the best performance, with its accuracy (acc), precision (pre), recall, F1 score (f1), and area under the curve (auc) being 85.87%, 86.85%, 86.39%, 86.62%, and 98.63%, respectively. In contrast, VGG showed the lowest performance among all the CNN models. As an early deep CNN model, VGG is simply composed of multiple convolutional and pooling layers stacked together, which results in a large number of parameters and, consequently, its relatively lower performance. Among the Transformer series, the Pyramid Transformer achieved the highest performance, with accuracy, precision, recall, F1 score, and AUC of 85.15%, 86.50%, 85.43%, 85.96%, and 98.60%, respectively. Nevertheless, overall, the performance of the Transformer series models was still slightly inferior to that of ResNet50. We speculate that this may be related to the size of the CRCTP dataset, as Transformer models typically adapt better to large-scale datasets, and may not easily demonstrate their potential advantages on relatively smaller datasets. Taking into account the above experimental results, we selected ResNet50 and the Pyramid Transformer (PVT-small) as the backbones for the dual-channel embedding network.

This experiment compared the impact of pre-training on the performance of the PVT model in few-shot classification tasks across three scenarios: same domain, near domain, and mixture domain. Meta-training was completed on the CRCTP dataset for all three scenarios, with meta-testing respectively performed on the CRCTP, NCTCRC, and LC25000 datasets, thereby simulating the aforementioned scenarios. The experimental results are shown in Table III.

Initially, we solely utilized PVT as the encoder component of a prototype network. After self-supervised pre-training, PVT achieved significant performance improvements in few-shot classification tasks across all three scenarios, with an average increase of 28.54%. This notable enhancement is primarily attributed to self-supervised learning conducting specific masking prediction tasks on a large volume of unlabeled data, enabling the model to learn rich feature representations without the need for annotations. This unsupervised pre-training approach provided PVT with a saturated feature space to alleviate its data-hungry characteristics, thereby reducing the risk of overfitting in few-shot data training tasks and enabling the successful application of transformers to few-shot classification tasks.

Subsequently, we constructed a dual-channel prototype network DCPN by combining PVT with ResNet and further explored the effect of pre-training within this architecture. Self-supervised pre-training also resulted in a performance increase in DCPN, with an average improvement of 3.48%. More crucially, in the absence of pre-training, DCPN exhibited a significant advantage over the prototype network using PVT alone, with a performance enhancement of 28.69%. We believe this gain stems from the inherent inductive bias of CNNs: local spatial continuity. When dealing with high-resolution medical images, local features may contain critical information about diseases or other medical conditions. Thus, CNNs provide the model with a more robust and fine-grained representation of pathological images. Such representation can compensate for the shortcomings of PVT, effectively mitigating its data-hungry issue.

The experiment aimed to compare the impact of features at different scales on the task of few-shot classification of pathological images, with the results summarized in Table IV. Here, ”global” and ”local” refer to features extracted by PVT and ResNet respectively, while ”mix” represents the concatenated features of global and local after dimensionality reduction through PCA. Furthermore, ”multi-scale” denotes the use of all three types of features in concert for the few-shot classification task of pathological images, as proposed in our DCPN architecture.

In the ”Same domain” scenario, the 1-shot, 5-shot, and 10-shot accuracy rates of ”local” features were 65.35%, 76.85%, and 78.39% respectively, which were notably superior to the accuracy rates of ”global” features, at 63.21%, 73.71%, and 75.89%. Further observations in the ”Near domain” and ”Mixture domain” scenarios revealed that ”local” features continued to exhibit better classification performance than ”global” features. These experimental results suggest that local features appear to have better discriminability in the few-shot classification tasks of pathological images, potentially because global features, which contain overarching information, may dilute the local characteristics beneficial for classification tasks in pathology, as demonstrated in Figure 6.

Additionally, examining the performance of ”mix” features, we find that by combining the principal components of local and global features, it outperforms the singular ”local” or ”global” features in all three scenarios. This confirms that integrating multiple scales of features can provide a more comprehensive image representation, thereby compensating for each other’s shortcomings. The feature visualization in Figure 6 also distinctly shows the differences and complementarity between PVT and ResNet features. The specific performance of ”multi-scale” features was the most outstanding, and we attribute this to several reasons: 1. The ”mix” feature, formed by the PCA-reduced concatenation of ”global” and ”local” features, may lose some discriminative features beneficial for classification; 2. The ”multi-scale” utilizes a soft voting strategy that can dynamically adjust the contribution of different scale features to the classification decision based on the confidence of the classifier outputs, thereby achieving optimal classification performance.

This experiment aimed to compare whether cosine similarity or Euclidean distance is more suitable for assessing the distance between query samples and prototype features within prototype-based few-shot classification methods for pathological images. The results in Table V indicate that across all three scenarios, employing Euclidean distance as the metric method outperforms cosine similarity. Particularly in the ”Same domain” scenario, the 5-shot classification task using Euclidean distance achieved an accuracy of 82.02%, significantly higher than the 75.58% using cosine similarity. Moreover, in the 10-shot task in the ”Mixture domain” scenario, using Euclidean distance achieved an accuracy of 65.01%, compared to 62.07% with cosine similarity, showing a clear advantage. This conclusion is consistent with research findings in the field of natural images [14].

Cosine similarity measures the similarity between two feature vectors by calculating the cosine of the angle between them, primarily reflecting the difference in vector direction rather than scale. In contrast, Euclidean distance involves not only the directional difference of vectors but also their scale difference, meaning the absolute numerical discrepancy between feature vectors. This suggests that Euclidean distance may provide more information when evaluating the differences between samples and prototype representations. It is noteworthy that when Bregman divergence is used as the metric, the prototype network can be seen as a process of probabilistic distribution estimation over the support set. However, cosine distance is not a Bregman divergence and hence does not possess the corresponding properties. This may explain why, in this study, the performance of squared Euclidean distance (a type of Bregman divergence) exceeded that of cosine distance.