CNN-Transformer Rectified Collaborative Learning for Medical Image Segmentation

Medical image segmentation (MIS) is a dense prediction task that classifies the pixels of organs or lesions in medical images (e.g., CT, MRI, US, endoscopy, etc.). Traditional MIS methods typically rely on hand-crafted features such as edge, shape, texture, etc, which suffer from a high risk of mis-segmentation for the inherent lack of high-level semantics.

With the development of deep learning technology, CNN-based MIS methods are first proposed and achieve impressive performance. Ronneberger et al. [7] proposed a symmetric U-shape network to aggregate multi-stage features with skip connections and progressively integrated them for high-resolution segmentation maps. Later, Zhou et al. [8] designed a series of nested, dense skip connections to reduce the semantic gaps between the feature maps and utilized deep supervision for multi-scale lesions. Jha et al. [23] introduced the residual connection [5], SE [24] and ASPP[25] to further improve the performance. However, limited by the receptive field in CNN, these methods are deficient in extracting global context information, leading to inaccurate localization. To this end, some Transformer-based MIS methods are proposed with self-attention mechanism applied for the long-range modeling. Karimi et al. [26] first presented a convolution-free segmentation model by forwarding flattened image representations to Transformer, whose outputs are then reorganized into 3D tensors to generate segmentation masks. Cao et al. [13] utilized the Swin Transformer to construct the encoder and decoder within the U-Net paradigm. Furthermore, Lin et al. [14] adopted a dual-scale Transformer encoder to attain coarse–fine-tuning features and employed the self-attention mechanism to establish global dependencies between them. Nonetheless, due to the local dilution in the self-attention process, methods based on pure transformer usually suffer from the coarse boundary in segmentation results. Besides, the excessive data dependency limits the potential of these methods for the insufficient medical data. Thus, some works couple both CNN and Transformer to design hybrid networks for MIS. Chen et al. [17] concatenated a Transformer-based encoder after the CNN-based encoder for global feature extraction. Zhang et al. [18] utilized shallow CNN-based encoder and Transformer-based segmentation network in parallel, and fused the features from two branches to jointly make predictions. Zhou et al. [27] constructed a CNN-Transformer hybrid backbone, nnFormer, with convolutional and Transformer blocks interleaved to take advantages of each other. Although these methods achieve better performance, the large model size and high computational cost prevent them from application. Different from the aforementioned methods, our CTRCL enables CNN-based model and Transformer-based model to learn from each other to improve the performance of each model.

Collaborative learning technology derives from the knowledge distillation [28] (KD) which transfers the knowledge from a cumbersome pre-trained teacher model to a compact student model for model compression. Differently, collaborative learning simultaneously exchanges the knowledge among a series of student models in an end-to-end training procedure to enhance the performance of each model.

Current collaborative learning methods can be primarily categorized into two types, mutual learning [29], [30], [31], [32] and ensemble learning [33], [34], [35], [36]. Mutual learning strategy is first time proposed by Zhang et al. [29] to make peer students learn from each other through the Kullback-Leibler (KL) divergence loss between each pair of student logits. Following this work, Chung et al. [37] introduced the mutual learning mechanism to feature maps using an adversarial training paradigm. More recently, Zhu et al. [38] proposed a bidirectional selective distillation strategy to transfer reliable knowledge between student models in both logit and feature spaces. Ensemble learning constructs a virtual teacher by assembling the outputs of all the students, which are then distilled back to foster each student. Lan et al. [33] constructed a multi-branch network and assembled the on-the-fly logit information from the branches to enhance the performance on the target network. Later, Guo et al. [39] extended the ensemble learning to the outputs of peer students with different architectures to generate high-quality labels for supervision. Furthermore, Kim et al. [40] introduced a feature fusion module which integrates the feature representations of sub-networks to construct the teacher classifier whose knowledge is delivered back to each sub-network. Unfortunately, current collaborative learning methods overlook the accuracy of the soft labels in the process of logit-wise knowledge transfer because neither a single student nor an ensemble teacher can completely reflect the distribution of the ground truth. Additionally, feature-wise collaborative learning between heterogeneous networks are not fully explored in previous studies. In contrast, we strive to ensure the accuracy of student soft labels with the help of the ground truth and to achieve effective feature-wise knowledge transfer between heterogeneous networks (CNN and Transformer) via the alignment of class-aware representations.

An overview of our CTRCL framework is depicted in Fig. 2, which consists of three components: a CNN-based student $f\left(\theta^{C}\right)$, a Transformer-based student $f\left(\theta^{T}\right)$, and the proposed collaborative learning strategies (RLCL and CFCL). Given an input image set $\mathcal{X}$ and its label set $\mathcal{Y}$, our objective is to enable $f\left(X;\theta^{C}\right)$ and $f\left(X;\theta^{T}\right)$ to learn collaboratively to assign the pixel-wise label $l\in 1,...,L$ in image $X\in\mathcal{X}\left(X\in\mathbb{R}^{H\times W\times D}\right)$ more accurately than the student itself, where $H$, $W$ and $D$ are the height, width and depth of $X$, $L$ is the number of categories. To achieve this goal, given a specific input $X$, we first attain the segmentation predictions ($P^{C}$ and $P^{T}$) and feature representations ($F^{C}$ and $F^{T}$) from the two students $f\left(X;\theta^{C}\right)$ and $f\left(X;\theta^{T}\right)$, formulated as:

Then the pixel-wise segmentation loss ($\mathcal{L}_{seg}$) is based on the cross entropy $\mathrm{CE\left(\cdot\right)}$ with the ground truth $Y$:

where ${\left(h,w\right)}$ indexes the spatial location for each pixel. Meanwhile, we propose a rectified logit-wise collaborative learning (RLCL) strategy and a class-aware feature-wise collaborative learning (CFCL) strategy to achieve bi-directional knowledge transfer between CNN-based and Transformer-based students in the logit and feature spaces, respectively. The details of RLCL and CFCL strategies are described as follows.

It is a fact that neither a CNN-based nor a Transformer-based student can completely mimic the ground truth, inevitably leading to mis-categorized pixels in their soft labels, which is adverse to the logit-wise collaborative learning between them. Consequently, we propose a rectified logit-wise collaborative learning (RLCL) strategy which adopts the ground truth to adaptively rectify the wrong regions in student soft labels for accurate logit-wise knowledge transfer between CNN-based and Transformer-based students.

As shown in Fig. 2 (b), the RLCL strategy introduces an adaptive rectification module (ARM) to rectify the segmentation predictions of both students ($P^{C}$ and $P^{T}$) and obtain rectified soft labels ($P^{C}_{r}$ and $P^{T}_{r}$) under the guidance of the ground truth $Y$, which can be formulated as:

The illustration of ARM is provided in Fig. 3. Given the segmentation prediction of a student as $P$, we first attain its pseudo label via the Argmax function and compare it with the ground truth $Y$ by the XOR operation to produce an error mask $M$ as follows:

Then we restore the one-hot label from the ground truth $Y$ and utilize the error mask $M$ to select the mis-segmented region $Y^{m}$ and $P^{m}$ from $Y$ and $P$, respectively. After that, $P^{m}$ is rectified to $P_{r}^{m}$ by the combination of $P^{m}$ and $Y^{m}$. The above process can be written as:

where $\mathrm{OHE}\left(\cdot\right)$ denotes the one-hot encoding and $\lambda$ is a matrix of dynamic weights for rectification demonstrated afterwards. Finally, we joint the rectified region $P_{r}^{m}$ and the correct region $P_{c}$ to generate the rectified soft label $P_{r}$, formulated as:

Now, we describe the computation of $\lambda$ and the process of rectification in detail. As a whole, the $\lambda$ consists of an alignment factor ${\lambda}^{a}$, a similarity-based decay factor ${\lambda}^{s}$ and a certainty-based decay factor ${\lambda}^{c}$, formulated as:

The distinct modeling approaches of CNN and Transformer result in intrinsic feature gaps between them. To this end, we propose a class-aware feature-wise collaborative learning (CFCL) strategy to achieve effective feature-wise knowledge transfer between heterogeneous CNN-based and Transformer-based students by granting their immediate features the similar capability of category perception.

As shown in Fig. 2 (a), the CFCL strategy introduces a category perception module (CPM) to extract the class-aware representations ($R^{C}$ and $R^{T}$) for the categories in ground truth $Y$ from immediate features ($F^{C}$ and $F^{T}$), which can be formulated as:

The structure of the CPM is illustrated in Fig. 4. Given a batch of input images $\left\{X^{n}\right\}^{N}_{n=1}$ where $N$ denotes the batch size, and their corresponding ground truths $\left\{Y^{n}\right\}^{N}_{n=1}$ with $K$ classes (including background) in total, let $F^{n}\in\mathbb{R}^{\hat{H}\times\hat{W}\times\hat{D}}$ be the feature map of a student network, where $\hat{H}$, $\hat{W}$ and $\hat{D}$ are the height, width and depth of $F^{n}$. The prototype of class $l\in\mathcal{K}$ (with $\left|\mathcal{K}\right|=K$) is generated via the masked average pooling and batch mean:

where the ground truth $Y^{n}$ is down-sampled to match the spatial size of the feature map $F^{n}$ with the nearest interpolation; $(\hat{h},\hat{w})$ indexes the spatial locations for each pixel; $\mathbbm{1}\left[\cdot\right]$ represents the indicator function that returns value 1 when the argument is true or 0 otherwise; $\epsilon$ is an infinitesimal to cope with the absence of some classes in a batch; and $N_{l}$ is the number of images that include the $l$th class.

Since the prototype actually denotes the channel-wise summary of a class, we can obtain the pixel-wise category perception by establishing the relationship between the feature map and the prototype. To achieve this objective, we adopt a non-parametric metric learning mechanism. Concretely, we first calculate the distances between the feature vector at each spatial location and all prototypes. Then we apply a softmax operation over the distances for the class-aware representation $R$. Let $\mathcal{P}^{n}=\left\{p_{l}\mid l\in\mathcal{K}^{n},\,\left|\mathcal{K}^{n}\right|=K^{n}\right\}$ be the prototypes for the $n$th sample where $K^{n}$ is the number of classes in the $n$th sample, for each $p_{i}\in\mathcal{P}$, we have:

where $d\left(\cdot\right)$ adopts the cosine distance, and multiplier $\alpha$ is a scaling factor fixed at 20 recommended by [41].

To achieve the class-aware feature-wise collaborative learning between CNN-based and Transformer-based students, we propose to minimize the loss function as follows:

Noteworthily, we perform the bi-directional CFCL on the output features of both encoder and decoder, i.e., $\left(F^{C}_{E},F^{T}_{E}\right)$ and $\left(F^{C}_{D},F^{T}_{D}\right)$, to exchange the information of high-level semantics as well as the low-level details between CNN student and Transformer student, respectively.

The overall loss functions of the CTRCL framework for CNN-based and Transformer-based students are given as:

where $\mathcal{L}_{seg}$ is the segmentation loss between the student predictions and the ground truth; ${\mathcal{L}}_{rlcl}$ is the rectified logit-wise collaborative learning loss between the student predictions; ${\mathcal{L}}_{cfcl}^{E}$ and ${\mathcal{L}}_{cfcl}^{D}$ denote the class-aware feature-wise collaborative learning loss between the student output features of encoder and decoder, i.e., $\left(F^{C}_{E},F^{T}_{E}\right)$ and $\left(F^{C}_{D},F^{T}_{D}\right)$ respectively; $\beta$, $\gamma_{1}$ and $\gamma_{2}$ are hyperparameters to trade off these loss terms. The optimization details are summarized in Algorithm 1. In test, we deploy both CNN-based and Transformer-based students because our target is to learn stronger CNN-based and Transformer-based models for MIS via the collaborative learning mechanism.

We adopt four evaluation metrics to evaluate the performance of various methods, including dice similarity coefficient (DSC), Jaccard index (JAC), Hausdorff distance (HSD) and mean absolute error (MAE). Dice similarity coefficient is a statistic used to gauge the similarity between two sets of data. Jaccard index is a measure used to evaluate the performance of an algorithm in the field of detecting objects within an image. Hausdorff distance is the maximum distance of a set to the nearest point in the other set. Mean absolute error is a quantitative measure of the average difference between all elements in the predicted mask and the ground truth. The higher value is better for dice similarity coefficient and Jaccard index while the lower is better for Hausdorff distance and mean absolute error.

We implement our proposed method on PyTorch framework and conduct all the experiments on a single NVIDIA RTX 4090 GPU with fixed random seeds. For CNN-based students, we choose the semantic FPN [45] as decoder with encoders of MobileNetV2 [46] and ResNet-50 [5]; for Transformer-based students, we utilize the efficient SegFormer [6] decoder with encoders of MiT-B0 and MiT-B2, which have comparable parameters with their CNN counterparts, respectively. Pre-trained weights on ImageNet-1k [47] are employed to initialize the backbone networks. We itemize the crop size (cs), batch size (bs), learning rate (lr), training epochs (ep), optimizer (opt) for each dataset: (1) Synapse Multi-organ dataset and ACDC dataset: cs=224$\times$224; bs=8; lr=0.0003; ep=200; opt=AdamW. (2) Kvasir dataset: cs=352$\times$352; bs=8; lr=0.00005; ep=300; opt=AdamW. We utilize the polynomial learning-rate strategy [48] to adjust the learning rate with the power of 0.9. We adopt the random flip and random rotation for data augmentation to avoid overfitting. For a more robust and fair comparison, we take the checkpoint from the last epoch and report the results on the testing sets.

We compare our proposed method with other 6 state-of-the-art collaborative learning methods, including DML [29], AFD [37], CTCL [38], KDCL [39], FFL [40], and AFID [44]. Among them, DML, AFD and CTCL are mutual learning based methods while KDCL, FFL and AFID are ensemble learning based methods. For comparison, we implement all the methods on two pairs: MobileNetV2 and MiT-B0; ResNet-50 and MiT-B2. Experimental results and analyses are as follows.

The proposed framework mainly consists of two key modules, namely rectified logit-wise collaborative learning (RLCL) and class-aware feature-wise collaborative learning (CFCL). In this section, we carry out a series of experiments with different settings to verify the performance of different components and key modules. We conduct our ablation studies on the Synapse Multi-organ dataset using ResNet-50 and MiT-B2, and choose DSC and HSD for evaluation.

In this section, we design a series of extensive experiments to demonstrate the generalization of our method. The results and analyses are as follows.