XNet v2: Fewer Limitations, Better Results and Greater Universality

Performance degradation with hardly HF information. As mentioned in [13], XNet is negatively impacted when images hardly have HF information. To intuitively illustrate this phenomenon, we compare the class activation map (CAM) [16] of HF encoder of XNet on CREMI [17] and ISIC-2017 [14]. From Figure 1, we can see that CREMI has rich HF information and HF encoder can better focus on these texture and edge details. In contrast, ISIC-2017 has less HF information, which prompts HF encoder to fail to extract recognizable information and locate specific segmentation objects.

Underutilization of raw image information. XNet uses LF and HF images generated by wavelet transform as input, and the raw images are not involved in training. Although LF and HF information can be fused into complete information in fusion module, the raw image may still contain useful but unappreciated information. Table I compares with and without raw images as input on ISIC-2017 [14] and we find that introducing raw images for the dual-branch further improves performance.

Insufficient Fusion. XNet only uses deep features for fusion. Shallow feature fusion and image-level fusion are also necessary. We introduce various fusions for XNet v2. Table VIII and Table XI in ablation studies of Section III-D demonstrate their effectiveness.

In view of the limitations of XNet, we propose XNet v2 and show its overview in Figure 2. XNet v2 consists of three sub-networks: main network $M$, LF network $L$ and HF network $H$. $M$, $L$ and $H$ are based on UNet [1] (3D UNet [18]). We use $L$ and $H$ to fuse with $M$ and use their respective shallow and deep features to construct L&M and H&M fusion modules, which enables $M$ to better absorb semantics and details. It also allows $L$ and $H$ to generate segmentation predictions with more perturbations.

Different from directly using LF and HF images generated by wavelet transform as input, XNet v2 performs image-level complementary fusion of LF and HF images, which further reduces limitations and improves universality (we discuss it in detail in Section II-C). The fusion results along with the raw images are fed into $L$, $H$ and $M$ to generate segmentation predictions for consistency learning.

XNet v2 uses LF and HF outputs to construct consistency loss with the output of $M$ respectively, which avoids the instability of training loss when LF or HF information is insufficient. To be specific, XNet v2 is optimized by minimizing supervised loss on labeled images and triple output complementary consistency loss on unlabeled images. The total loss $L_{total}$ is defined as:

where $L_{sup}$ is supervised loss, $L_{unsup}$ is unsupervised loss, i.e., triple output complementary consistency loss, $\lambda$ is a weight to control the balance between $L_{sup}$ and $L_{unsup}$. Same as [13], $\lambda$ increases linearly with training epochs, $\lambda=\lambda_{max}\ast epoch/max\_epoch$. We compare the performance of different $\lambda_{max}$ in ablation studies of Section III-D.

The supervised loss $L_{sup}$ is defined as:

where $p_{i}^{M}$, $p_{i}^{L}$ and $p_{i}^{H}$ represent segmentation predictions of $M$, $L$ and $H$ for the $i$-$th$ image, respectively. $y_{i}$ represents ground truth of the $i$-$th$ image. The unsupervised loss $L_{unsup}$ is defined as:

same as [13], $L_{unsup}^{M,L}(\cdot)$ and $L_{unsup}^{M,H}(\cdot)$ are achieved by CPS [19] loss:

where $\hat{p}_{i}^{M}$, $\hat{p}_{i}^{L}$ and $\hat{p}_{i}^{H}$ represent pseudo-labels generated by $p_{i}^{M}$, $p_{i}^{L}$ and $p_{i}^{H}$, respectively. As in [13], all losses are dice loss [2].

For inference process, we use the segmentation predictions of $M$ as the final result.

Different from [13], after using wavelet transform to generate LF image $I_{L}$ and HF image $I_{H}$, we fuse them in different ratios to generate complementary image $x^{L}$ and $x^{H}$. $x^{L}$ and $x^{H}$ are defined as:

where $\alpha$ and $\beta$ are the weights of $I_{H}$ and $I_{L}$, respectively. We can see that the input of XNet is a special case when $\alpha=\beta=0$ while our definition is a more general expression. Figure 1 intuitively compares $x^{L}$, $x^{H}$ with different $\alpha$, $\beta$.

Simple but Effective. This strategy is simple but achieves image-level information fusion. More importantly, it solves the limitation of XNet not working with less HF information. To be specific, when hardly have HF information, i.e., $I_{H}\approx 0$:

$x^{H}$ degenerates into a perturbation form of $x^{L}$, which can be regarded as consistent learning of raw images with two different LF perturbations. It effectively overcomes the failure to extract features when HF information is scarce.

We set $\alpha$ and $\beta$ to change randomly within the range $[a,b]$ during training stage, which increases the diversity and randomness of training samples to further improve training quality. We compare different range combinations of $\alpha$, $\beta$ and demonstrate the effectiveness of image-level fusion in ablation studies of Section III-D.

We use fusion module to transfer feature-level complementary information between $L$ and $M$, $H$ and $M$. Taking L&M fusion module as an example, we describe its structure. We use $E_{n}^{M}$ and $E_{n}^{L}$ to represent the $n$-$th$ layer features of $M$ and $L$, respectively. The fusion between $E_{n}^{M}$ and $E_{n}^{L}$ is shown in Figure 3. $E_{n}^{M}$ and $E_{n}^{L}$ perform channel concatenation to acquire features with twice the number of channels. Then we use 3×3 convolution to fuse features and concatenate the fused features to the decoders of $M$ and $L$.

For $M$ and $L$, we use deep ($3rd$ and $4th$) features for fusion. For $M$ and $H$, we use shallow ($1st$ and $2nd$) features for fusion. The design of two fusion modules is asymmetric, which is also equivalent to introducing feature-level perturbations into the model.

We evaluate our model on three 2D datasets (GlaS [20], CREMI [17] and ISIC-2017 [14]) and two 3D datasets (P-CT [15] and LiTS [21]). Their preprocessing is the same as [13, 12].

Following [13], We use Jaccard index (Jaccard), Dice coefficient (Dice), average surface distance (ASD) and 95th percentile Hausdorff distance (95HD) as evaluation metrics.

We implement our model using PyTorch. Training and inference of all models are preformed on four NVIDIA GeForce RTX3090 GPUs. For 2D datasets (GlaS, CREMI and ISIC-2017), the initial learning rate is set at 0.8. For 3D datasets (P-CT and LiTS), the initial learning rate is set at 0.05. Other experimental setups (such as momentum, training epoch, batch size, training size, etc.) are the same as [13].

Semi-Supervision. We compare XNet v2 extensively with 2D and 3D models on semi-supervised segmentation, including UAMT [22], URPC [23], CT [24], MC-Net+ [25], etc. From Table II and Table III, we can see that XNet v2 significantly outperforms previous state-of-the-art models in both 2D and 3D. Furthermore, because of the introduction of image-level complementary fusion and the effective utilization of raw images, XNet v2 has more competitive performance than XNet and is capable of handling scenarios where XNet cannot work (such as on the ISIC-2017 and P-CT datasets), which addresses the limitation of XNet in handling insufficient HF information.

Fully-Supervision. The comparison results are shown in Table IV. As previous experiments, XNet v2 still shows superior performance compared to XNet.

To verify effectiveness of each component, we perform the following ablation studies in semi-supervised learning.

Comparison of Range Combinations of $\alpha$ and $\beta$. As shown in Equation 5, different range combinations of $\alpha$ and $\beta$ produce different LF and HF complementary fusion images. To determine the optimal range combination , we conduct comparative experiments on GlaS. We set 3 value ranges for $\alpha$ and $\beta$ to generate 9 combinations. Table V shows the results for GlaS and we find that larger $\alpha$ and $\beta$ achieves better performance. According to the analysis in Section II-C, this may be because larger $\alpha$ and $\beta$ alleviate the performance degradation of insufficient LF or HF information.

Comparison of the Trade-off Weight $\lambda_{max}$. The comparison results of different $\lambda_{max}$ on five datasets are shown in Table VI. We find that for relatively easy datasets (GlaS, ISIC-2017 and P-CT), $\lambda$ should increase faster (i.e., $\lambda_{max}$ should large) to highlight the role of many unlabeled images to prevent overfitting. For more difficult datasets (CREMI and LiTS), $\lambda$ should change smoothly (i.e., $\lambda_{max}$ is small), so that the model can better use the labeled images in the early training stage and further improve from unlabeled images in the later training stage.

Effectiveness of Wavelet Perturbation. We compare the wavelet perturbation with other common perturbations in Table VII, including Gaussian noise, network initialization, image smoothing and sharpening. We find that wavelet perturbation achieved better results. We also note that smoothing and sharpening can also enhance LF semantics and HF details but have a negative impact. Furthermore, we also apply various perturbations to MT [8] and acquire consistent conclusions.

Effectiveness of Image-Level Complementary Fusion. We compare the performance of different image-level fusion strategies, including without fusion ($\alpha\&\beta=0$), single-sided fusion ($\alpha|\beta=0$) and complementary fusion ($\alpha\&\beta\neq 0$). From Table VIII, we can see that single-sided fusion hardly has positive effect. It may be because only using single-sided fusion cannot effectively transfer complementary information to the other branch and thus affects the calculation of consistency loss. In contrast, complementary fusion can improve performance by a large margin, because it realizes the mutual complementation of missing frequency information.

Effectiveness of Raw Images. As mentioned in Section II-A, the information of raw images is also crucial for segmentation. In Table IX, we show the performance improvement of XNet v2 by introducing raw images. Furthermore, introducing additional branches for raw images can further improve performance, so we design additional main network $M$ for raw images in XNet v2.

Comparison of Model Size and Computational Cost. To illustrate that the performance improvement comes from well-designed components rather than the additional parameters brought by multiple networks. We compare the performance of semi-supervised models with the similar scale on GlaS and the results are shown in  Table X. We find that the increase in the number of parameters (Params) and multiply-accumulate operations (MACs) cannot bring positive effects to these semi-supervised models. Furthermore, as in [13], we reduce the number of channels of XNet v2 to half and quarter to generate XNet v2^- and XNet v2^–. These lightweight models still have superior performance than lightweight XNet with similar scale (XNet^- and XNet^–). The above experiments strongly prove that the performance improvement comes from various designs rather than the increase in model size and computational cost.

Effectiveness of Components. To demonstrate the improvement of different components, we conduct step-by-step ablation studies on GlaS and the results are shown in Table XI. Using raw images as input and training the semi-supervised model based on three independent UNet, we achieve a baseline performance of 78.80% in Jaccard. Using LF and HF complementary fusion images as input improves the baseline by 1.86% and 2.71% in Jaccard, respectively. Using them together further improves the baseline to 81.65% in Jaccard. Introducing L&M and H&M fusion modules improves the baseline by 3.88% and 3.31% in Jaccard, respectively. By introducing all components, we finally improve the baseline to 83.17% in Jaccard.