E2ENet: Dynamic Sparse Feature Fusion for Accurate and Efficient 3D Medical Image Segmentation

To Reviewer SYnH

We would like to thank the reviewer for your thoughtful and detailed comments. We are glad that you appreciate the efficiency improvements and robust validation presented in our paper. We address your comments below.

Unclear Backbone Network: The backbone network of E2ENet is not clearly defined. Figure 2 labels the left part as an "efficient backbone," but it is ambiguous whether this refers to an EfficientNet-based backbone or simply a group of CNN layers. This lack of clarity can confuse readers and detract from the paper's overall comprehensibility.

Thank you for pointing out the description of the backbone network, which indeed helps to enhance the readers' understanding and reduce ambiguity. We will provide this information below and add it to our paper.

The efficient backbone network consists of several levels as in Figure 2, each comprising two consecutive blocks. Each block includes our proposed efficient Restricted Depth-Shift 3D Convolutional layer, followed by instance normalization and ReLU activation (termed conv–norm–relu). After each level, the downsampling is performed using a strided convolution operation in the second block of that level (the convolution in the second block of the new resolution has a stride >1). Table: Backbone Network – Each row describes level $i$, with input resolution $H^i$, $W^i$, $D^i$, and output channels $C^i$, given an input size of 128×128×128.

Kernel Size Generalization: Question: Why was the 3x3x3 kernel size not included in the ablation study results for the BraTS dataset in Table 4? Suggestion: It would be beneficial to include the results for the 3x3x3 kernel size in the ablation study for the BraTS dataset.

Thank you for your valuable suggestion. Indeed, a 3x3x3 kernel size is helpful in the ablation study for the BraTS dataset as well. This will aid in understanding the generalizability of the DSFF mechanism across different datasets. Here, we update these results as follows and include them in Table 4 of the paper:

We find that, compared to kernel sizes of 3x3x3, E2ENet with kernel sizes of 1x3x3 combined with a depth shift (3rd and 4th row) maintains or even improves segmentation accuracy, whether with DSFF or without DSFF. This further demonstrates that our proposed efficient Restricted Depth-Shift 3D Convolutional layer, which utilizes a 1x3x3 kernel with restricted depth shift, is equivalent to a 3x3x3 kernel in terms of segmentation accuracy. Moreover, it offers significant savings in computational and memory resources, as observed in the AMOS-CT dataset in Table 3.

Comparison with Recent Lightweight Networks: Question: Have you considered comparing E2ENet with other recent lightweight network structures specifically designed for medical image analysis? Suggestion: Including comparisons with recent lightweight networks would strengthen the paper by providing a more comprehensive evaluation of E2ENet's performance.

Thank you for your suggestion. We agree that comparisons with recent lightweight networks would strengthen the paper. We have included another recent efficient lightweight network, UNETR++ [1], as our baseline. UNETR++ offers both high-quality segmentation accuracy and efficiency. From the results below, we can see that E2ENet achieves better mDice scores. This further verifies the efficiency and accuracy of our proposed model for 3D medical segmentation.

Moreover, we conducted an extended ablation study by integrating our proposed DSFF into UNETR++[1], referred to as UNETR++ w/ DSFF in the table below. Specifically, we introduced dynamic sparsity into UNETR++ by sparsifying its connections: less important activated connections are removed, while the same number of deactivated connections are randomly reactivated during training. The results show that parameters can be further reduced while maintaining stable mDice performance. This further demonstrates the effectiveness of our proposed DSFF, showing that it can be easily integrated into recent lightweight networks to potentially enhance efficiency while maintaining mDice performance.

[1] Abdelrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang and Fahad Shahbaz Khan. UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation. IEEE Transactions on Medical Imaging, 2024.

We thank you again for the time and effort you've taken to participate in the review of our paper. If you have further questions and concerns, we are more than happy to discuss with you.

Dear Reviewer SYnH,

We sincerely thank you for your constructive comments and support!

Best,

Authors

To Reviewer JUDh

Thank you very much for taking the time to review our paper and for your helpful comments. We provide detailed responses to your constructive feedback below.

For Section 3.2 Table 3, it seems that for the different shift size, the mDice score does not vary much from each other. What are authors’ insights about this?

Thank you for your enthusiasm for our paper and for sharing your concerns with us. Indeed, we observe a marginal difference in mDice, with a decrease of 0.4 when the shift size magnitude is increased from 1 to 3. However, for mNSD, there is a significant performance drop from 82.3 to 81.6. This is because the mDice score primarily measures the overlap between the predicted segmentation and the ground truth, whereas mNSD focuses on the distance between the surfaces (boundaries). This demonstrates that shift size is more impactful on boundary alignment.

Furthermore, in additional experiments where we extended the shift size magnitude to 7, we observed a further decrease in performance by 2.5 for mDice and 4.8 for mNSD, compared to a shift size magnitude of 1. The detailed results are shown below. In summary, the shift operation is beneficial, resulting in a 1.9 improvement in mDice compared to not using the shift operation. However, our experiment results indicate that the magnitude should not be set too large. We will explain more below.

How do authors justify in Section 2.3 that the shift magnitude would have a negative impact on the effectiveness of the shift operation?

In Section 2.3, we propose the Restricted Depth-Shift operation, which aims to capture depth-wise information while maintaining a 2D computation cost. However, increasing the shift size means considering more depth-wise information at the expense of channel-wise information, leading to an insufficient representation of channels, as discussed in Section 3. Additionally, a large shift size causes a loss of local spatial relationships, which are crucial for segmentation. This results in a blurring effect that reduces the precision of boundary alignment, particularly affecting metrics like mNSD, which rely heavily on accurate boundary information.

We would like to thank you again for your time and effort. If you have further questions, we are more than happy to discuss them with you.

To Reviewer xwUb

We sincerely thank you for your time and effort in reviewing our paper. We are glad that you find our work interesting and novel. We address your comments below.

Tables 3 and 4 don’t provide consistent results. It seems that the combination that works for CT does not optimally work for MRIs.

For the mDice score, in Table 3, the shift operation in E2ENet with and without DSFF combination scores 90.1 and 90.2, respectively. In Table 4, with and without DSFF scores 81.2 and 81.0, respectively. These results rank among the top in these tables and show marginal differences in performance.

However, with the DSFF combination, while it does not significantly improve the mDice scores, the computational cost in terms of parameters is notably reduced from 23.9M to 11.2M, a 2-fold reduction for both Table 3 and Table 4.

• Performance: The mDice scores show marginal differences, indicating that the shift operation combined with DSFF maintains comparable performance to the shift operation without DSFF.
• Efficiency: The incorporation of DSFF into E2ENet results in a substantial reduction in the number of parameters and FLOPs, achieving a significant computational cost reduction.
• Conclusion: As discussed in Section 3.2, our main conclusion is that dynamic sparse feature fusion significantly reduces computational resources while maintaining high performance without any significant degradation.

Furthermore, in Appendix B.7.3, we provide additional verification on the statistical significance of the designed modules: DSFF and depth-shift operation.

Can the proposed method accelerate training speed in terms of time to achieve SOTA accuracies? Could the authors provide information on the inference times?

Thank you for your enthusiasm for our paper and for sharing your concerns with us. Given the relatively restricted support for sparse operations in current off-the-shelf commodity GPUs and TPUs without sparsity-aware accelerators, we did not attempt to achieve practical speedup during training. Instead, we chose to implement our models with binary masks in our work. As mentioned extensively in the conclusion and appendix, the promising benefits of dynamic sparsity presented in this study have not yet translated into actual speedup. Accelerating training time will be a focus for our next work.

Although not the focus of our current work, it would be interesting for future work to examine the speedup results of sparse operation during training, using such specialized hardware accelerators, as we see much improvement room of promise here. For example, at high unstructured sparsity levels, XNNPACK (https://github.com/google/XNNPACK) has already shown significant speedups over dense baselines on smartphone processors.

For the inference time:

Although the support for unstructured sparsity on GPUs remains relatively limited, its practical relevance has been widely demonstrated on non-GPU hardware, such as CPUs and customized accelerators. For instance, FPGA accelerators for an unstructured sparse RNN achieved high acceleration and energy efficiency compared to commercial CPUs and GPUs by maximizing the use of embedded multiply resources on the FPGA. Another notable success is DeepSparse (https://github.com/neuralmagic/deepsparse), which successfully deploys large-scale BERT-level sparse models on modern Intel CPUs, achieving a 10× model size compression with less than 1% accuracy drop, a 10× CPU-inference speedup with less than 2% accuracy drop, and a 29× CPU-inference speedup with less than 7.5% accuracy drop.

Inspired by these advancements, we adopted an approach based on DeepSparse. We conducted experiments with patches of images sized 32×32×32 as input, comparing the CPU wall-clock timings for online inference between our proposed E2ENet and nnUNet on an Intel Xeon Platinum 8360Y CPU with 18 cores. We acknowledge that, while our proposed models with sparsity do achieve speedups in practical inference, they are not as pronounced as those observed with BERT-level sparse models. This is primarily due to the nature of segmentation and 3D convolution operations. However, this presents a promising avenue for our future work.

Missing discussion on limitations.

We appreciate the reviewer pointing this out. I will make it more clear in the final version. Our proposed model leverages unstructured dynamic sparsity; however, due to the relatively restricted support for unstructured sparsity in current off-the-shelf commodity GPUs, we did not attempt to achieve practical speedup during training on GPU. This limitation presents a potential direction for our future work.

We thank you again for the time and effort you've taken to participate in the review of our paper. If you have further questions, we are more than happy to discuss with you.