An efficient computational framework for gastrointestinal disorder prediction using attention-based transfer learning

Benchmark dataset

To develop our method, we used the KVASIR dataset (Pogorelov et al., 2017). The dataset contains 8,000 images, divided into three groups with eight sub-groups. The group Phatological Findings (disease detected) has three sub-groups: Esophagitis, Polyps, Ulcerative Colitis. The group Polyp Removal is specified by two sub-groups: Dyed and Lifted Polyps and Dyed Resection Margins. The group Anatomical Landmarks (normal colon) has three sub-groups: Z-line, Pylorus, and Cecum. Each sub-group comprises 1,000 images whose resolution is 224 ×224 ×3 pixels. We prepared three sets of data: a training set, a validation set, and an independent test set with 7120, 80, and 800 images, respectively. Table 1 gives information on the number of images in the training, validation, and independent test sets.

Transfer learning

Transfer learning (TL) is a technique in the field of ML that involves applying knowledge gained from solving one problem to a different but related problem (Torrey & Shavlik, 2010). This approach is particularly useful when there is a scarcity of data available for the target task. Essentially, a pre-trained model, developed for a task with abundant data, is repurposed and fine-tuned to improve performance on the new task. This process not only saves time and resources but also enhances learning efficiency, as the model does not need to learn from scratch. TL has found wide applications across various domains, such as image recognition and natural language processing. By leveraging pre-trained networks, it enables more accurate and efficient training of models on specific, often limited datasets. This approach is particularly beneficial in scenarios where collecting large amounts of labeled data is impractical or too costly. In this study, we used four pre-trained networks, including ConvNeXt (Liu et al., 2022), DenseNet-161 (Huang et al., 2016), EfficientNetV2 (Tan & Le, 2021), and ResNet-18 (He et al., 2015), as the main networks for feature extraction. Additional linear layers are added after the main networks for fine-tuning to learn distinct features of images. To enhance learning efficiency, a generic attention mechanism is integrated to help the model focus on essential regions (Vaswani et al., 2017).

Attention mechanism

To enhance the focus on specific regions of images in fine-tuned layers, we used Scaled Dot-Product attention, a simplified and commonly used attention mechanism (Vaswani et al., 2017). This mechanism operates on the principle of calculating attention by scaling the dot products of Queries and Keys. Essentially, it involves dividing the dot products by the square root of the dimension of the Keys, which helps in stabilizing the gradients during training. This approach enables the model to efficiently weigh the significance of different parts of the input data. By focusing on relevant parts while processing sequences, this attention facilitates a more nuanced understanding and representation of the input. It is especially crucial in tasks that involve processing sequences, like language translation or text summarization, where context and the relationships between different elements in a sequence are key. The computation of Scaled Dot-Product attention is: (1)${\displaystyle Attention\left(Q,K,V\right)=softmax\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V,}$ where Q, K, V denote Query, Key, and Value vectors, respectively, and d_k is a dimension vector.

Proposed model

Model training

Figure 1 displays the proposed architecture of the CAD system. Initially, images with a resolution of 224 ×224 ×3 enter a pre-trained network. This network is responsible for extracting features from the images, which are then used in the subsequent layer. Utilizing pre-trained networks for image feature extraction significantly saves time, resources, and effort. Besides enhancing performance, these networks can mitigate the risk of overfitting, especially in scenarios with small datasets, due to their robustness and comprehensive initial training. During training, the weights of the pre-trained network were fixed. The size of the extracted feature vectors varies depending on the type of pre-trained network used. These feature vectors are then fed into a fully connected (FC) layer for fine-tuning before passing through an Attention block. The attention block contains three parallel linear layers for the transformation of feature vectors into corresponding Q, K, and V vectors. The attention output vectors are passed through a linear layer and then activated by the Softmax function to return prediction probabilities. All models were trained over 30 epochs and optimized using the SGD optimizer (Ruder, 2016) with a learning rate of 0.001. During the training process, only the linear layers (after the main network) were fine-tuned, while the weights of the main network remained unchanged. The loss function used was cross entropy.

Model evaluation

To find the suitable attention mechanism, we applied three different types of attention, including Dot Product, Additive, and Scaled-Dot Product. After training, these three models were examined using the validation set. The results show that the model employing Scaled-Dot Product attention works more effectively compared to those employing other types of attention (Table 2). Hence, we chose Scaled-Dot Product attention for developing our model.

Table 3 indicates the experimental results on the performance of CAD systems developed across four types of pre-trained networks. The results show that the CAD systems developed with ConvNeXt obtain a higher AUC_ROC value, followed by the DenseNet-161-based, ResNet-18-based, and EfficientNetV2-based systems (Fig. 2). Although the DenseNet-161-based CAD system achieves an AUC_PR value of over 0.92, higher than other CAD systems, it is still lower than the ConvNeXt+Attention-based system, which is enhanced by the attention block. In terms of other metrics, the ConvNeXt-based system is still dominant over the others. However, after being enhanced by the attention mechanism, the CAD system shows superior performance compared to ConvNeXt-based models as well as other systems. The addition of the attention mechanism evidently helps to improve the prediction efficiency.

Model explainability

To investigate the attention-based learning of our model, we visualized the attention maps of eight classes of samples to highlight specific focusing regions across different classes, emphasizing these regions’ importance in distinguishing one class from another. The visualization of these attention maps contributes to understanding the model’s decision-making process, providing insights into its attention mechanism. Figure 3 illustrates the attention maps returned by the model for eight classes. These visualized maps show that the model focuses on each class’s distinct regions. The attention maps of the normal cecum and z-line classes indicate that their right-handed regions are less focused compared to those of the other classes. The attention maps for the esophagitis, polyps, and ulcerative colitis classes reveal that their left-handed regions have darker color bands, which means these regions are more focused than other regions. The darker color bands in the attention maps of disease samples demonstrate that our model effectively learns to distinguish them from normal samples.

In addition to visualizing attention maps, we explored the distribution of samples from eight classes in the training data and test data using t-distributed stochastic neighbor embedding (t-SNE). The output vectors of the final FC were used as feature vectors for t-SNE visualizations (Fig. 4). These t-SNE plots show that after passing through the attention layer, all classes are almost separately distributed except for the polyps and dyed lifted polyps classes. The overlapping population between these two classes can be explained by their similar shapes found in some locations. Overall, the model has effectively learned to separate all classes to perform classification tasks.

Stability evaluation

To examine the robustness and stability of the model, we conducted the experiments five times, each with a different random sampling seed for dividing the dataset into the training, validation, and test sets. This approach was designed to evaluate the variation in the model’s performance across these distinct partitionings. The results consistently revealed very minor fluctuations in performance metrics, such as accuracy, precision, and recall, across the different trials (Table 4). This consistency confirms that our proposed method is not only effective but also exhibits a high level of stability, making it reliable for further application and suitable for reproduction in various settings. Such robustness is crucial for ensuring that the model can handle diverse and unpredictable real-world data effectively.

Limitations

ConvNeXt (Liu et al., 2022), DenseNet-161 (Huang et al., 2016), EfficientNetV2 (Tan & Le, 2021), and ResNet-18 (He et al., 2015) are prominent DL architectures, each with unique features and varying performance characteristics. ConvNeXt evolves from the traditional CNN design, similar to ResNet-18, but incorporates insights from Transformers (Vaswani et al., 2017), leading to improved image classification performance. Its strength lies in its ability to balance the robustness of CNNs with the efficiency of Transformers, though it might not be as extensively tested as more established models. DenseNet-161, on the other hand, stands out with its dense connectivity pattern, where each layer receives input from all preceding layers, significantly reducing the vanishing gradient problem and enhancing feature propagation. However, this can lead to increased model complexity and computational demands. EfficientNetV2, an advanced version of EfficientNet, excels in balancing speed and accuracy, particularly in mobile and real-world applications, due to its scalable architecture and focus on efficiency. It boasts improved training speed and higher accuracy but may require more careful tuning for specific tasks. ResNet-18, the simplest among these, uses residual connections to enable the training of deeper networks without the gradient vanishing issue, making it efficient for basic tasks but potentially less effective for more complex applications compared to its more advanced counterparts. Each of these models represents a significant approach to neural network design, with pros and cons. Under the scope of this study, although the ConvNeXt+Attention-based CAD system yields higher performance compared to the others, its limitations need to be considered and addressed in further studies.