[論文レビュー] 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ベースバックボーンを提案する。並列なMLPブランチによる三平面オフセット学習と、構造的再パラメータ化にインspiredされた設計を統合することで、4つの公開データセットで最先端の性能を達成し、SOTAのトランスフォーマーと比較してパラメータ数を削減した。

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深さ方向の可変的畳み込みの可能性の検証。
• 学習可能なチャネルに依存するオフセットを用いて、特徴マップのチャネル間での自己適応的空間アグリゲーションの向上。
• SOTAのViTを上回る性能とパラメータ効率性を実現する新しいCNNベースの基礎モデルの確立。

提案手法

• エンコーダブロックに深さ方向の可変的畳み込みを用いた、U-Netに類似した3D DeformUX-Netの導入により、有効な受容 field を拡大。
• すべてのチャネルにわたる自己適応的空間アグリゲーションを可能にするために、1×1×1畳み込みを用いた並列ブランチで三平面可変オフセットを生成。
• 構造的再パラメータ化の原則を応用し、MLPを介して可変オフセットを計算することで、特徴表現とモデル効率性を向上。
• MLPを用いた線形スケーリングにより空間特徴のアグリゲーションを強化し、深さ方向スケーリングを上回る性能を達成。
• 多方向の空間的依存性を捉えるために三平面オフセット計算を適用し、解剖的構造によって性能にばらつきが生じる。
• 4つの公開3D医療画像データセット上で、SOTA手法と直接比較可能な教師ありのエンドツーエンド訓練を実施。

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ベースのバックボーンは、パラメータ数を減らしながらも、SOTAのビジョントランスフォーマーと同等の性能を達成できるか？

主な発見

• KiTS（腎臓、腫瘍、嚢胞）において、DeformUX-Netは平均Diceスコア0.720を達成し、既存のSOTAモデルを上回った。
• MSD膵臓データセットでは、平均Diceが0.717に達し、臓器セグメンテーションタスクにおける強力な性能を示した。
• AMOSにおける多臓器セグメンテーションでは、DeformUX-Netは平均Diceが0.902に達し、SOTAのViTおよびCNNを上回った。
• 肝臓血管セグメンテーションでは、平均Diceが0.671に達し、微小で複雑な構造に対しても頑健であることを示した。
• MLPベースの三平面オフセットの使用は、すべてのデータセットで一貫して性能向上をもたらし、特に血管セグメンテーションにおいてx-z平面が最も効果的であった。
• トランスフォーマー基盤のモデルと比較してパラメータ数を削減しながら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