[论文解读] DeformUX-Net: Exploring a 3D Foundation Backbone for Medical Image Segmentation with Depthwise Deformable Convolution

Ho Hin Lee, Quan Liu+5 more|arXiv (Cornell University)|Sep 30, 2023

Advanced Neural Network ApplicationsComputer Science被引用 3

一句话总结

DeformUX-Net 提出了一种用于医学图像分割的3D基础主干网络，通过深度可分离可变形卷积统一了长距离依赖、自适应空间聚合与计算效率。通过引入并行多层感知机分支进行三平面偏移学习，并采用结构重参数化启发式设计，该方法在四个公开数据集上实现了最先进性能，且参数量少于当前最先进模型的视觉Transformer（ViTs）。

ABSTRACT

The application of 3D ViTs to medical image segmentation has seen remarkable strides, somewhat overshadowing the budding advancements in Convolutional Neural Network (CNN)-based models. Large kernel depthwise convolution has emerged as a promising technique, showcasing capabilities akin to hierarchical transformers and facilitating an expansive effective receptive field (ERF) vital for dense predictions. Despite this, existing core operators, ranging from global-local attention to large kernel convolution, exhibit inherent trade-offs and limitations (e.g., global-local range trade-off, aggregating attentional features). We hypothesize that deformable convolution can be an exploratory alternative to combine all advantages from the previous operators, providing long-range dependency, adaptive spatial aggregation and computational efficiency as a foundation backbone. In this work, we introduce 3D DeformUX-Net, a pioneering volumetric CNN model that adeptly navigates the shortcomings traditionally associated with ViTs and large kernel convolution. Specifically, we revisit volumetric deformable convolution in depth-wise setting to adapt long-range dependency with computational efficiency. Inspired by the concepts of structural re-parameterization for convolution kernel weights, we further generate the deformable tri-planar offsets by adapting a parallel branch (starting from $1 imes1 imes1$ convolution), providing adaptive spatial aggregation across all channels. Our empirical evaluations reveal that the 3D DeformUX-Net consistently outperforms existing state-of-the-art ViTs and large kernel convolution models across four challenging public datasets, spanning various scales from organs (KiTS: 0.680 to 0.720, MSD Pancreas: 0.676 to 0.717, AMOS: 0.871 to 0.902) to vessels (e.g., MSD hepatic vessels: 0.635 to 0.671) in mean Dice.

研究动机与目标

• 解决3D医学图像分割中全局-局部注意力机制与大感受野卷积之间的权衡问题。
• 开发一种计算高效、自适应的主干网络，能够捕捉长距离依赖关系，且不依赖于自注意力机制。
• 探索3D深度可分离可变形卷积作为体积分割基础算子的可行性。
• 通过可学习的、通道感知的偏移量，改进特征通道间的自适应空间聚合。
• 建立一种新型基于CNN的基础模型，其性能超越当前最先进视觉Transformer（ViTs），在参数效率方面更具优势。

提出的方法

• 提出3D DeformUX-Net，一种类似U-Net的架构，在编码器模块中使用深度可分离可变形卷积，以扩展有效感受野。
• 采用并行分支，通过1×1×1卷积生成三平面可变形偏移量，实现在所有通道上的自适应空间聚合。
• 借鉴结构重参数化原理，通过多层感知机（MLP）计算可变形偏移量，提升特征表示能力与模型效率。
• 采用线性缩放策略结合MLP增强空间特征聚合，性能优于深度缩放策略。
• 应用三平面偏移计算以捕捉多方向的空间依赖关系，其性能因解剖结构而异。
• 在四个公开的3D医学影像数据集上以端到端监督方式训练模型，以便与当前最先进方法进行直接比较。

Figure 1: This figure compares our proposed block design with representative 3D medical image segmentation designs. We leverage depth-wise deformable convolution in parallel with a multi-layer perceptron (MLP), which generates the tri-planar offset to adapt long-range dependency and adaptive spatial

实验结果

研究问题

• RQ1 3D深度可分离可变形卷积是否能有效替代3D医学分割中的全局-局部自注意力机制与大感受野卷积？
• RQ2 三平面偏移计算在不同解剖结构（如器官和血管）上的性能影响如何？
• RQ3 采用可学习的、并行分支设计进行偏移量生成，是否能提升自适应空间聚合能力，同时减少参数量？
• RQ4 基于MLP的偏移量计算是否在体积分割任务中优于固定核或标准可变形卷积？
• RQ5 基于CNN的主干网络若结合可变形操作，是否能实现超越视觉Transformer（ViTs）的SOTA性能，且参数量更少？

主要发现

• 在KiTS数据集（肾脏、肿瘤、囊肿）上，DeformUX-Net实现了0.720的平均Dice分数，优于现有SOTA模型。
• 在MSD胰腺数据集上，其平均Dice达到0.717，展现出在器官分割任务中的强大性能。
• 在AMOS多器官分割数据集上，DeformUX-Net的平均Dice达到0.902，超越当前SOTA的ViTs与CNN模型。
• 在肝血管分割任务中，其平均Dice为0.671，表现出对小而复杂结构的鲁棒性。
• MLP驱动的三平面偏移量在所有数据集上均一致提升性能，其中x-z平面在血管分割中效果最佳。
• 该模型在参数量少于基于Transformer的模型前提下实现了SOTA性能，证实了其计算效率与可扩展性。

Figure 2: Overview of the deformable convolution mechanisms. Deformable convolutions introduce an adaptable spatial sampling capability that transcends the rigid bounds of conventional $3\times 3\times 3$ regions, achieved with the deformable offsets (light green arrows). Such offsets can demonstrat