CFII-Net: Explicit Class Embeddings and Feature Maps Through Iterative Interaction for Boosting Medical Image Segmentation

1 Introduction

Medical image segmentation is a fundamental task in computer vision, critical for diagnosis and preoperative planning. Recently, CNN-based and Transformer-based models have achieved significant progress in this domain. A classical method U-Net [Ronneberger et al., 2015] uses an encoderdecoder structure with skip connections for dense predictions. Following this technical route, its derived variants in U-Net with advanced network block techniques [Zhou et al., 2018;

∗Corresponding Author

Xiao et al., 2018; Oktay et al., 2018; Jin et al., 2019]. Vision transformers [Dosovitskiy et al., 2020] highly focus on learning the relationships between different patch tokens, which is an effective method for global context modeling. For example, TransUNet [Chen et al., 2021] and SwinUNet [Liu et al., 2021] both combine CNNs and transformers through serial fusion, obtaining a better trade-off in accuracy and efficiency. TransFuse [Zhang et al., 2021b] introduces the BiFusion module to parallelly combine global dependencies and spatial details. However, these methods focus on obtaining better pixel representations while neglecting the structural prior knowledge of the target classes implicit in the feature maps. This can lead to semantic inconsistencies within the same class and confusion between different classes.

Recently, the mask classification segmentation paradigm has aroused the interest of many researchers in natural image segmentation tasks. This approach typically involves feeding class embeddings and feature maps into a transformer decoder to facilitate interaction, ultimately predicting a set of binary masks, each linked to a specific class. For example, SegVit [Zhang et al., 2022] leverages class embeddings and attention maps to identify local patches with higher similarity, improving inference efficiency. Maskformer [Cheng et al., 2021] and Mask2former [Cheng et al., 2022] use the class embeddings to generate class predictions per segment, addressing semantic and instance-level segmentation tasks. Although class embeddings in the methods above are linked to category information, they are typically initialized randomly, meaning they are implicit and reliant on continuous training to learn the structural information of the categories.

In general, medical images possess significant class structural prior knowledge (e.g., organ anatomy and spatial relationships) compared to natural images, enabling the extraction of explicit and meaningful class embeddings to assist segmentation. Intuitively, shallow features capture finegrained details, while deeper features represent abstract semantic information. In other words, each stage's features thus provide a comprehensive depiction of category-specific characteristics. Consequently, class embeddings derived from these feature maps inherit similar traits, guiding feature map regions belonging to the same category to cluster cohesively.

In light of this promising attempt, we try to extract explicit class embeddings with category structural prior in medical images and consider a simple method for predicting segmentation maps using Explicit Class Embeddings and Feature Maps through Iterative Interaction, called CFII-Net, which consists of two parts: Explicit Class Embedding Generator (ECEG) and Iterative Interactor. Specifically, the ECEG overcomes the limitations of class-agnostic feature maps by retaining part of the original features and activating the class attributes of other features, which effectively improves the quality of class embeddings later. Subsequently, by aggregating the masks of the same class, we generate prototype class vectors to obtain explicit class embeddings. To achieve more refined class embeddings and segmentation maps, we introduce iterative interactor based on an attention mechanism. Employing this mechanism, class embeddings continuously extract information from features to improve class discriminability, and features continuously explore class prior knowledge from embeddings to enhance class relations and semantics. CFII-Net produces better results than SOTA methods with a significantly lower parameters cost (as shown in Figure 1) on synapse dataset.

Our main contributions can be summarized as follows:

• We utilize category information of feature maps to generate explicit class embeddings with structural prior knowledge to guide medical image segmentation, instead of blindly relying on random initialization.
• We build a novel iterative interaction framework CFIINet, where the high-quality class embeddings are obtained by ECEG module, and pixel-to-class relationships between feature maps and class embeddings are explored by the Iterative Interactor, progressively refining the results.
• Extensive experimental results on four public medical datasets validate that CFII-Net outperforms state-of-theart methods in accuracy and parameters count.

2 RelatedWork

2.1 CNN and Transformer for Medical Image Segmentation

Hybrid architectures based on convolutional neural network and transformer have shown excellent performance in medical image segmentation. Scaleformer [Huang et al., 2022] from a scale-wise perspective to improve the segmentation quality even for small objects. UDTransNet [Wang et al., 2023] explores the skip connection between encoder-decoder levels by exploring channel and spatial attention mechanisms. DTMFormer [Wang et al., 2024] addresses the potential attention collapse in hybrid frameworks by proposing a plugand-play module with dynamic token merging. EMCAD [Rahman et al., 2024] proposes a new and efficient multiscale convolutional attention decoder to obtain a better tradeoff in accuracy and efficiency.

2.2 Class Embeddings for Image Segmentation

DETR [Carion et al., 2020] builds an end-to-end detector by introducing learnable object queries in the transformer decoder to correspond to each object instance, which has given a lot of food for thought to image segmentation. Segmenter

[Strudel et al., 2021] uses a mask transformer to process image patches and class embeddings jointly and applies several self-attention layers to produce prediction maps. Maskformer [Cheng et al., 2021] first realizes the importance of class embeddings and uses it to predict N class labels and N corresponding mask embeddings. Then the Mask2former [Cheng et al., 2022] further advances this by introducing a new transformer decoder and employing an optimization strategy to reduce training memory. OMG-Seg [Li et al., 2024] follows the Mask2Former architecture and proposes a new encoder-decoder to handle all segmentation tasks efficiently. SegViTv2 [Zhang et al., 2024] transfers the similarity mapping between a set of learnable class embeddings and feature maps to the ATM modules to obtain segmentation masks. ECENet [Liu et al., 2023] improves segmentation performance using roughly predicted segmentation masks for generating class embeddings.

However, since medical images contain richer category structural prior knowledge compared to natural images, randomly initialized class embeddings lead to the loss of valuable prior knowledge. Based on this, we attempt to extract meaningful explicit class embeddings from the feature maps, which are then used to iteratively interact with the features to boost medical image segmentation.

3 Method

The proposed CFII-Net framework tailored for explicit class embedding for medical image segmentation consists of two modules, i.e., 1) Explicit Class Embedding Generator (ECEG), and 2) Iterative Interactor. As shown in Figure 2, given an input image of size H×W, we first use a backbone (such as ResUNet) to generate multi-stage features F. Subsequently, we use chunk (a simple tensor operation) on F along the channel dimension to form dual branches multistage features F1 and F2, which are used for obtaining class embeddings and refined feature maps respectively. These are defined as F1,i ∈ RHi×Wi×Ci and F2,i ∈ RHi×Wi×Ci , ∀i = 1, 2, 3, 4, 5, 6, where Hi = H / 2i-1 and Wi = W / 2i-1 , represents the i-th stage feature whose scale is 1/2i-1 of the input image, Ci is the channel dimensions of stage i. For clarity of representation, we replace F2 with F throughout the paper. One branch F1 is used to generate the corresponding prototype explicit class embeddings E by the designed ECEG module. Then, the class embeddings and the other branch F are fed into multi-stages of Interactors to interact and update them iteratively. Finally, we send feature maps of each stage after interaction to the downstream block for prediction. Each module of our CFII-Net is described in detail below.

3.1 Explicit Class Embedding Generator (ECEG)

Initial multi-layer feature maps are often class agnostic, so directly mapping the feature maps to class embeddings is imprecise. To this end, we introduce a simple but effective method to obtain high-quality class semantic embeddings, dubbed Explicit Class Embedding Generator (ECEG), based on the decomposition of prediction masks.

Since the diversity and intrinsicity in features are important for obtaining high-quality category embeddings [Woo et al., 2018; Han et al., 2020], as observed in Figure 2, we randomly choose a branch of multi-stage features F1 and contend that the feature maps F1,i ∈ RHi×Wi×Ci produced in each stage contain intrinsic features. Hence, we employ linear transformations ϕ(·) implemented by 1 × 1 convolutional without activation to extract these features in each stage.

To enhance the model's response to the focal points, we introduce Channel Spacial Attention (CSA) to generate attention feature maps information sequentially in both channel and spatial dimensions. These are then multiplied with the original input feature maps for adaptive feature correction, resulting in an enhanced feature maps.

where CSA sequentially infer a 1D channel attention map Mc ∈ RCi×1×1 and a 2D spatial attention map Ms ∈ R1×Hi×Wi . ⊗ denotes element-wise multiplication and [·] denotes concatenation. Finally, the intrinsic properties branch F ′ and category properties branch F ′′ are concated together and projected to the initial shape. In this way, we obtain enhanced features X for helping to get high-quality semantic embeddings.

Generally, convolution and attention-based operations do not explicitly decouple the category cues, so extracting category information is non-trivial. Mask classification is the assignment of a mask to each class, ideally the region in which the class is located is 1 and the background is 0. Based on this, it is logical to consider using prediction masks as the most natural raw material for extracting explicitly defined class embeddings. As shown in Figure 2, after obtaining the enhanced features Xi ∈ RHi×Wi×Ci , we first apply linear transformations to predict a coarse segmentation probability map Yi, which can be written as follows:

where ϕ1 and ϕ2 are linear transformations implemented by 1 × 1 convolutional layers without activation, N equals to the number of classes. Then we disassemble Yi by category, with each class corresponding to a class mask, and use the sigmoid function to highlight the category features and generate the category gate maps Yn i ∈ RHi×Wi , n ∈ {1, . . . ,N}. We then apply both Yn i and category gate map inversion Yni for weighted fusion of the enhanced features Xi.

The key insights are two folds. First, sharing the same gates can better highlight category regions. Second, the use of a subtracted gate technique supplements the missing details in the nonsalient parts. Such process is shown as follows:

Finally, we aggregate the representations of all pixels, compute their average as the prototype vector for a category, and concatenate all the category prototype vectors as explicit class embeddings E.

where Pooling is the adaptive average pooling. A set of category prototype vectors is integrated by a prediction mask obtained from the enhanced feature maps, which in turn generates meaningful class embeddings. Rather than randomly initializing the class embeddings blindly, our approach explicitly extracts prior knowledge of the category structure hidden in the feature maps, which is important in the subsequent segmentation process.

3.2 Interactive Iterator

After obtaining the multi-stage explicit class embeddings E, we feed it together with the multi-stage features F into the Interactive Iterator. As shown in Figure 2, we use SelfInteractive Transforme (SIT) at the deeper stage of the network and Cross-Interactive Transforme (CIT) at the shallower stage, respectively. Meanwhile, we employ Class Embeddings Updater (CEU) and Feature Enhancement Updater (FEU) to update the class embeddings and feature images progressively, which obtains more fine-grained explicit class embeddings and segmentation maps.

Self-Interactive Transformer (SIT)

As illustrated in Figure 3(a), we reshape the image features Fi ∈ RHi×Wi×Ci into a sequence of patche tokens ˜Fi ∈ RHiWi×Ci , where i ∈ {3, 4, 5, 6}, then concatenate them with the class embeddings Ei ∈ RN×Ci to obtain Gi ∈ R(N+HiWi)×Ci , and perform self-attention on Gi to ensure that the image features can explicitly recognize the class dependencies from the class embeddings. Specifically, linear transformations are applied to Gi to form query (Q),

where ϕα∈{q,k,v} are the linear transformations, √D serves as a scaling factor while D equals to the dimension of key. Subsequently, we pass through the feed-forward networks (FFN) and split the class embeddings and patche tokens along the channel dimension. And then reshape the patch tokens into feature maps,

where E′i ∈ RN×Ci and F′i ∈ RHi×Wi×Ci are the interacted class embeddings and interacted features obtained by reversing the order of the concatenated sequence, and are used in the iterative update in the next part.

Cross-Interactive Transformer (CIT)

The problem of quadratic complexity becomes increasingly obvious when the feature map resolution is large, so updating the feature maps sensibly is non-trivial. We aim to use class embeddings as category structural prior knowledge to help feature maps obtain clearer segmentation results, so the feature maps fed into the decoder are of utmost importance. We can effectively solve this problem by using crossattention. As illustrated in Figure 3(b), it reshapes the image features Fi ∈ RHi×Wi×Ci into a sequence of patche tokens ˜Fi ∈ RHiWi×Ci , where i ∈ {1, 2}, which are then used as a query (Q) with the class embedding as key (K) and value (V ), and the rest of the operations are similar to those in SIT.

It is worth noting that, in addition to obtaining unidirectional (E→F) interacted of the feature maps F′i through the reshape operation, we also obtain interacted class embeddings E′i through matrix multiplication by the mid-product attention map associated with the query.

Compared with the conventional self-attention with the complexity of O(CH2W 2 ), the cost of cross-attention is decreased and can be summarized as O(2CNHW), where C is the number of channels, H,W are the height and width of the feature maps, respectively, and N is the number of segmentation categories. Even though multiple categories N are considered, the cost of cross-attention is relatively small compared to self-attention (2N ≪HW).

Class Embeddings Updater (CEU)

At each stage, there are unique prototype class embeddings with varying prior knowledge of class structure. Therefore, we need to continuously enrich the class embeddings.

As illustrated in Figure 3(c), after obtaining the interacted class embeddings E′i ∈ RN×Ci of the i-th layer, we refresh the previous class embeddings Ei-1 ∈ RN×Ci-1 of the (i - 1)-th layer through CEU. We first concatenate E′i and Ei-1 along the channel dimension for each category, and then apply LayerNorm (LN) [Ba et al., 2016] to standardize the features of each sample:

where chunk(·) indicates separation along the channel dimension. {·} represents concatenation by category, so n ∈ {1, . . . ,N}, N is the number of categories. Since the class features from different sources may have different distributions, this operation helps the model distinguish between different categories, thereby improving the class distinguishability. Inspired by [Zhang et al., 2021a], we use a gating mechanism to refresh the previous class embeddings further. Finally, we obtain the updated class embeddings.

ˆEi-1 = Ψ(σ(E) ⊙LN(chunk(ϕ4(E′i ))) +Ei-1) (16) where σ is the sigmoid function and Ψ is a fully connected (FC) layer followed by LayerNorm. This iterative update gradually repeats on multi-stage features, producing highquality class embeddings.

Feature Enhancement Updater (FEU)

Multi-scale information is particularly important for segmentation tasks. Therefore, we need to enhance the multi-scale feature information continuously. As illustrated in Figure 3(d), where the feature maps of each level are compressed to the same channel dimension by two 1 × 1 convolutional layers before entering the next level. Given an interacted feature maps F′i and a previous-level feature maps Fi-1, both of which have the same channel dimension, we upsample F′i to the same size as Fi-1 by a bilinear interpolation layer. Then, inspired by [Li et al., 2019], we attempt to adapt the mechanism of automatic selection between different maps. Specifically, we first integrate the information of two branch feature maps.

Then, we embed the global information through simple maximum pooling and average pooling and generate channel and spatial feature descriptors. To allow the interaction of different feature descriptors, we use a simple fully connected layer and a sigmoid activation function to obtain compact features Aspa and Acpa for accurate and adaptive selection guidance.

where Pmax(·) and Pavg(·) are the maximum and average pooling. Finally, F′i and Fi-1 are matched with spatial feature descriptor and channel feature descriptor, respectively, to maximally preserve semantic information, and we get the updated outcomes ˆFi-1 :

In brief, CEU and FEU units are reused in multi-stages, resulting in highly informative class embeddings and feature maps. The class distinguishability of the class embeddings is progressively enhanced, and valuable semantic details in the feature maps are extracted while reducing irrelevant contextual information.

Decoder and Loss

Finally, the interacted multi-level feature maps from Interactive Iterator are sent to a series of decoder blocks, including two 3×3 convolutions and a skip connection. In the training phase, we use the combined Cross-Entropy loss LCE and Dice loss LDice as the primary loss. Moreover, we introduce an auxiliary Euclidean distance loss LEU [Li et al., 2022] to maximize the distance between the interacted class embeddings, which can further enhance class distinguishability.

where LEU is balanced by λEU, which is further shown in ablation experiments.

4 Experiments

4.1 Datasets

Synapse (9 classes) consists of 30 abdominal CT scans. Following [Chen et al., 2021], we split 18 cases for training and 12 cases for testing. We report the Dice Coefficient (Dice) and 95% Hausdorff Distance (HD95) on 9 different organs. ACDC (4 classes) contains 100 MRI scans involving three organs. Consistent with [You et al., 2022], we present the Dice results using a random split of 70 training cases, and 30 testing cases. MoNuSeg (2 classes) contains 44 images, with 30 images for training, and 14 for testing. Following [Kumar et al., 2017], we report Dice and Intersection over Union (IoU) as the evaluation metrics. GlaS (2 classes) [Sirinukunwattana et al., 2017] consists of 85 training samples and 80 testing samples. we report Dice as the evaluation metrics.

4.2 Implementation Details

We implement our model with PyTorch on a single NVIDIA 4090 GPU card with 24 GB of memory. To avoid overfitting, we also perform two data augmentations including random rotation and flipping. We use ResUNet as the backbone and do not use any pre-trained weights to train the proposed CFIINet. For all four datasets, we set the input resolution to 224 × 224 and generally train for 600 epochs.

For Synapse and ACDC, we use the SGD optimizer to train our model with a batch size of 8, where the initial learning rates are set to 0.006 and 0.04, respectively, with momentum of 0.9 and weight decay of 1e - 4. For MoNuSeg and GlaS, we train our model using the Adam optimizer, where the initial learning rates are set to 0.001 and 0.004, respectively, with a batch size of 4, and use CosineAnnealingWarmRestarts as a scheduler, with a maximum number of iterations of 10, and a minimum learning rate of 1e - 4.

4.3 Comparison with the State-of-the-Arts

Quantitative Comparison

We compare it with 12 state-of-the-art (SOTA) networks on the Synapse dataset, including V-Net [Milletari et al., 2016], U-Net, AttUNet, ViT, TransUNet, SwinUNet, UCTransNet [Wang et al., 2022a], MTUnet [Wang et al., 2022b], AFTerUNet [Yan et al., 2022], MissFormer [Huang et al., 2021], ScaleFormer, and EMCAD. As shown in Table 1, although CFII-Net achieves sub-optimal results in HD95 (slightly inferior to EMCAD), it outperforms all others in the Dice metric at 84.58%. In particular, compared to ScaleFormer which only aims at extracting feature representations at different scales, we utilize the paradigm of explicit class embedding to guide the segmentation process, which maintains a better Dice metric (Dice: +1.72%) and a drastic reduction in model parameters (Params: -88.09M). Table 2 shows the Dice scores on the ACDC dataset, where CFII-Net obtains the highest average Dice score of 92.37% and outperforms other methods in all three organ segmentations.

In addition to multi-category organ segmentation, we also validate the effectiveness of CFII-Net on two binary medical image datasets. The experimental results on GlaS and MoNuSeg are reported in Table 3, where our method achieves significant advantages over others.

Qualitative Comparison

Figure 5 presents qualitative results for two binary segmentation datasets, demonstrating that CFII-Net produces outputs closer to the ground truth with fewer false segmentations. Figure 4 showcases results for multi-category segmentation datasets, highlighting the limitations of existing methods. For the ACDC dataset, EMCAD under-segments RV organs, while ScaleFormer, MissFormer, and TransUNet undersegment Myo organs. In the Synapse dataset, other methods exhibit various segmentation errors: EMCAD confuses KL and KR due to insufficient class guidance, ScaleFormer and MissFormer misidentify the background as pancreas, UCTransNet segments the background as left kidney, TransUNet misidentifies the background as liver and stomach, SwinUNet under-segments the aorta, and UNet mistakenly segments the background as gallbladder and left kidney, while confusing the right kidney with the left kidney. In contrast, CFII-Net achieves superior segmentation by leveraging explicit class embeddings enriched with structural prior knowledge. These results underscore the potential of explicit class embeddings as a promising paradigm for medical image segmentation.

4.4 Ablation Studies

Components in CFII-Net

Table 4 highlights the contributions of different components in CFII-Net. The introduction of explicit class embeddings (W/CLS), a core element of our method, significantly improves segmentation performance (Dice: +2.75%) compared to the baseline. Moreover, incorporating FEU (Dice:+1.97%) and CEU (Dice:+1.90%) modules to iteratively update feature maps and class embeddings further enhances performance with negligible additional parameters. Unlike blind randomly initialized class embeddings, we introduce explicit class embeddings with prior knowledge of the category structure to aid segmentation, which largely improves segmentation performance (Dice:+2.75%). These demonstrate that guided by explicitly defined class embeddings, feature maps can yield refined class cues, thus improving segmentation performance.

Different Backbones

We integrate the class embeddings with three other representative medical image segmentation architectures, including CNN-based U-Net, transformer-based TransUNet, and SwinUNet, i.e., using our CFII-Net except for the backbone and decoder parts. As shown in Table 5, we find that CFII-Net achieves consistent improvements in both Dice and HD95, and also slightly reduces the model parameters.

Loss Coefficient

Finally, we evaluate the impact of different loss coefficients on performance using the Synapse dataset. The choice of loss coefficient λEU is entirely result-driven. Table 6 shows that our CFII-Net achieves an optimal 84.58% Dice when λEU = 0.3. Therefore, we use this setting in all experiments.

5 Conclusion

We present CFII-Net which uses explicit class embeddings to boost the medical image segmentation. Compared with the traditional random initialization, explicit class embedding possesses category structural prior knowledge, so regions of same category in the feature maps are more inclined to be clustered together and different categories are more distinguishable under the guidance of explicitly defined class embedding. In addition, we gradually refine feature maps and class embeddings during the iterative process, which improves the performance of the network. CFII-Net achieves state-of-the-art performance on four publicly available datasets.

Acknowledgments

This work was supported in part by the National Nature Science Foundation of China (62273150, 62171323) and the Science and Technology Commission of Shanghai Municipality under Grant 22DZ2229004.