QTSeg: A Query Token-Based Dual-Mix Attention Framework with Multi-Level Feature Distribution for Medical Image Segmentation

To effectively capture the inter-channel relationships within the features, we introduce the CAB module, which performs self-attention on the feature channels, as depicted in Fig. 4c. The CAB module is built upon the attention mechanism, where $F$ is considered the number of tokens and $H\times W$ represents the feature dimension of the tokens. The formulation of the query $Q\in\mathbb{R}^{F\times HW}$, key $K\in\mathbb{R}^{F\times HW}$, and value $V\in\mathbb{R}^{F\times HW}$ of the CAB block can be expressed as follows:

where $\mathcal{S}_{skip}=Q=K=V=Flatten(\mathcal{S})$ and $Flatten(\cdot)$ is the module that flattens features on the feature size axis. A skip connection is added to Equation (1) to mitigate the vanishing gradient issue and facilitate effective learning by reutilizing the previous features. Similarly, a Spatial Attention Block (SAB) can be achieved through self-attention by treating the channel as a token dimension and the feature size as the number of tokens. However, this approach results in an attention matrix $\mathbb{R}^{HW\times HW}$, which then requires multiplication by the value matrix to compute the attention. This process can be computationally expensive and slow, making it less suitable for high-resolution images and real-world applications. Fortunately, this issue can be mitigated by incorporating SAM [18], which is much faster than SAB and does not depend on the feature size. The idea of SAM is depicted in Fig. 4d, which can be expressed as follows:

where $\sigma$ is sigmoid activation, $f^{3\times 3}$ represents a convolution operation with filter size of $3\times 3$. By leveraging SAM for channel attention, DMA can mitigate the weak long-range dependencies in CNN architectures at the feature level.

Another key feature of DMA is its capability to align the query tokens with the image features using cross-attention [19]. We propose the utilization of cross-attention for each image feature in query tokens, namely CTFA, enabling the integration of current image features into the query tokens. CTFA employs attention to integrate information from image features to query tokens, where the query tokens are considered as $Q\in\mathbb{R}^{F\times N}$, and the image features from CAB or SAM is represented as $K\in\mathbb{R}^{F\times HW}$ and $V\in\mathbb{R}^{F\times HW}$ in formula (3).

where $Q^{\prime}$ = $Q^{T}W^{q}$, $K^{\prime}$ = $K^{T}W^{k}$, $V^{\prime}$ = $V^{T}W^{v}$ are the projection of $Q,K,V$ with learnable matrix $W^{q},W^{k},W^{v}$ and $d_{k}$ is the feature dimension of $K^{\prime}$. The final query tokens, which encompass feature attention information, are obtained by summing the query tokens and the outputs of the two CTFA modules applied to CAB and SAM as follows:

After improving the representation of feature embedding and query tokens through DMA, we apply the Cross-Feature Token Attention (CFTA), which is similar to CTFA, to apply token attention to features. In CTFA, the query tokens will be considered as $K\in\mathbb{R}^{F\times N}$ and $V\in\mathbb{R}^{F\times N}$, while the feature embeddings are represented as $Q\in\mathbb{R}^{F\times HW}=\mathcal{S}_{CAB}+\mathcal{S}_{SAM}$. The process from Equation (1) to (4) is iterated with $h$ blocks before proceeding to another CTFA for query token alignment. Finally, the output of the DMAD will be fed into the upsampling and MLP blocks to project the feature dimensions to the previous stage.

The ISIC2016 [33] dataset specifically focuses on enhancing melanoma diagnosis through the utilization of high-quality, human-validated datasets comprising skin lesion images. The challenge provided official training and testing datasets featuring 900 dermoscopic lesion images for training and 379 for testing. These images were accompanied by ground truth masks labeled by experts and saved in binary mask format. The BKAI-IGH NeoPolyp [34] is the polyp segmentation dataset released by the BKAI Research Center at Hanoi University of Science and Technology in collaboration with the Institute of Gastroenterology and Hepatology (IGH) in Vietnam. This dataset comprises 1,200 endoscopic images, with a division of 1,000 images for training and 200 images for testing purposes. While an official split of 1,000 training and 200 testing samples exists, the official test ground truth is not public due to the ongoing challenge. Thus, we partition the training set into five folds to evaluate the model’s performance comprehensively. To ensure the reproducibility of results and facilitate future comparisons, we first sort all samples based on their sample names. Subsequently, these samples were segmented into K folds by array slicing. In our experimental setup, we executed the models across five folds and subsequently reported the averaged metrics for a comprehensive evaluation.

The BUSI [35] dataset comprises a collection of breast ultrasound images obtained from 600 female patients aged between 25 and 75 years. This dataset includes 780 images categorized into three classes: normal, benign, and malignant. As the normal images do not exhibit any lesions, a new dataset was obtained by excluding these normal images. The refined dataset exclusively contains cases classified as benign (437 images) and malignant (210 images), which were then partitioned using the K-fold strategy.

The DSB2018 [36] was published in 2018, focusing on nucleus segmentation across many imaging experiments. This diverse collection includes images under various conditions, showcasing differences in cell type, magnification, and imaging modalities such as brightfield and fluorescence. Each image is annotated with a unique identifier, linking it to corresponding segmentations, which are provided in training datasets. We also partition the training set into five folds to evaluate the model’s performance comprehensively.

The FIVES [37] aimed to develop and assess artificial intelligence models for segmenting blood vessels in fundus images. This dataset comprises a total of 800 samples with high-quality color and resolution. The primary goal is to enable precise, efficient, and automated segmentation of retinal vessels, a crucial step in diagnosing and managing a range of clinical conditions. The official training and testing datasets include 600 fundus images for training and 200 fundus images for testing.

To mitigate model overfitting during training, we implemented the augmentation process described in MISSFormer [7]. Our QTSeg model underwent training for 350 epochs on each dataset, using a batch size of 32 and the AdamW optimizer with a learning rate set to 0.001 and a weight decay of 0.0001. A scheduler was employed to adjust the learning rate, reducing it by a factor of 0.5 every 50 epochs. For retinal vessel segmentation specifically, a batch size of 4, the SGD optimizer with an initial learning rate of 0.01, and a poly scheduler strategy were employed. Input images were resized to 512 x 512 for lesion, breast cancer, and cell segmentation and 1024 x 1024 for retinal vessel segmentation. Following resizing, these images were normalized to a value range of [0,1] using min-max normalization. To maintain fairness in comparisons, no post-processing was applied to the output results. All experiments were conducted on an NVIDIA GeForce RTX 3090 GPU running on the Debian 12 system. For evaluation metrics, we utilized standard metrics, including Mean Absolute Error (MAE), Accuracy (Acc), Intersection-over-Union (IoU), and Dice Similarity Coefficient (Dice).

Table I presents the performance of QTSeg on the lesion, breast cancer, and cell segmentation datasets, outperforming recent methods across all evaluation metrics. Specifically, our model achieves impressive scores of 92.45% and 86.92% in Dice and IoU metrics, respectively, for skin lesion segmentation. For the BKAI-IGH NeoPolyp dataset, our model also attained the highest scores in Dice and IoU metrics, showcasing its effectiveness. Furthermore, QTSeg demonstrates great potential by delivering outstanding results with considerably lower FLOPs and parameter counts compared to other high-accuracy methods. It is worth noting that QTSeg’s parameter count is four times lower than that of H2Former [6] while achieving superior performance across all metrics. While EGE-UNet [38] and MALUNet [39] have smaller parameter sizes, they encounter challenges in converging on the BKAI-IGH NeoPolyp dataset due to inherent design limitations and model parameter constraints. Our approach stands out for achieving precise lesion and polyp segmentation with minimal error compared to alternative methods.

On the other hand, as depicted in Fig. 5, the visualization of our model’s predictions compared to other methods highlights QTSeg’s superior accuracy, with the generated mask closely resembling the ground truth mask. Fig. 6 further demonstrates our model’s clean and accurate segmentation mask, emphasizing its effectiveness in medical image segmentation.

The lesion segmentation experiments above mainly consist of single segmentation objects, which do not fully showcase the performance of QTSeg in multiple object segmentation within medical images. To further evaluate the performance of QTSeg in a different segmentation scenario, we conducted an experiment using the DSB2018 dataset. As depicted in Table I, QTSeg demonstrates superior performance compared to other approaches, with a significant lead of 1.97% in the IoU metric over the second-best method. This result underscores the effectiveness of our proposed method in general multiple object segmentation and the specific task of cell segmentation.

Fig. 7 visually illustrates the segmentation performance comparison between our proposed method and other high-scoring approaches. Notably, our method successfully classifies a background area that is similar to the cell object, as evidenced by the first red bounding box, a task that other methods struggle with. This accurate classification in the first red bounding box highlights the effectiveness of our DMA in integrating global information into the decoder, providing crucial contextual information for mask decoding. Moreover, even in architectures such as TranUNet and Swin-Unet, known for handling long-range dependencies, the first red bounding box is identified as a cell, further showcasing the efficiency of our DMA module above these approaches.

On the other hand, while CNN-based models excel in capturing information in the second red bounding box, transformer-based models seem to face challenges in this area, as depicted in Fig. 7e-h. Nonetheless, our model exhibits these limitations, as indicated by the yellow box, where improvements are needed compared to the labeled image. This constraint arises from the presence of blur and small objects in the image, requiring further enhancement in the future.

As depicted in Table I, QTSeg demonstrates superior performance across all metrics when evaluated using a five-fold assessment. In particular, our QTSeg model achieves remarkable results with 82.09% Dice and 73.85% IoU, outperforming FSCA-Net by 0.84% and 0.90%, respectively. In Fig. 8, a failure case study on the BUSI dataset showcases QTSeg’s accurate predictions compared to other methods. Despite challenging features in the sample, QTSeg excels in segmenting a portion of the sample accurately, unlike models like U-Net [8] and MISSFormer [7] that struggle with the segmentation task. While approaches like MHorUNet [40] and TransUNet [14] can predict the target area, our QTSeg model outperforms them by generating predictions that closely match the ground truth with significantly less error. This visualization effectively illustrates how the DMA module leverages global information to enhance its predictions, especially when dealing with similar features extracted by the encoder model. In Fig. 8a, the input image lacks distinct features for differentiation, leading to segmentation areas that closely resemble the background, as evident in the misclassification shown in Fig. 8g of a CNN-based method. In contrast, employing a similar approach to the CNN-based method, QTSeg successfully discriminates between the background and foreground regions by incorporating information from other areas. This exemplifies how the DMA module effectively addresses the challenges related to local-global information integration in image segmentation tasks.

It is crucial to assess our method’s performance on high-resolution images and various segmentation tasks. Table II provides a comparison of our approach with other recent methods using Acc, Dice, and IoU metrics. Once more, QTSeg stands out by achieving the highest scores across all metrics: 98.93% for Acc, 90.37% for Dice, and 84.26% for IoU. It is evident that QTSeg excels in accurately segmenting complex structures and nuances present in high-resolution images. The consistently high scores in Acc, Dice, and IoU metrics signify the robustness and reliability of our approach in capturing fine details and boundaries within the images. Furthermore, the superior performance on vessel segmentation tasks highlights the versatility and effectiveness of QTSeg across different segmentation scenarios. Overall, these results underscore the potential of our method as a leading solution in image segmentation tasks.

A deep visualization analysis is presented in Fig. 9 for a deeper understanding of our model performance. The red and yellow dash bounding boxes show the improvement and the shortcomings of the proposed method, respectively. The QTSeg can distinguish the blood vessel from the background and does not misclassify the vessel according to the second red bounding box in the figure. Besides, the first red bounding box, which has features similar to those of the blood vessel, has a good classification compared to other methods. However, our method still gains inferior performance on some blurred blood vessels or similar vessels, as highlighted in the yellow bounding boxes. The design of the DMA module successfully addresses the challenges posed by local and global information constraints in both CNN and Transformer architectures. The visual representation in the figure illustrates how our method minimizes mask errors while also being computationally efficient and parameter-light compared to alternative approaches. Moreover, the masks produced by QTSeg demonstrate a remarkable ability to differentiate between similar features like colors and shapes, enabling accurate inferences based on the blood vessel structures. This capability is exemplified by the distinct delineation of blood vessels and lesions within the first red bounding box in the figure. These observations emphasize the significance of integrating transformer mechanisms with CNN feature extraction, enabling our model to leverage global information effectively.