[論文レビュー] 3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation

Ho Hin Lee, Shunxing Bao+2 more|arXiv (Cornell University)|Sep 29, 2022

Medical Imaging and Analysis被引用数 99

ひとこと要約

この論文は、 volumetric medical image segmentation の階層的なトランスフォーマー挙動を模倣する大カーネル深さ方向畳み込みを用いた純粋なConvNetである 3D UX-Net を提案し、パラメータ数を抑えつつ、公開データセットで最先端の結果を達成する。

ABSTRACT

The recent 3D medical ViTs (e.g., SwinUNETR) achieve the state-of-the-art performances on several 3D volumetric data benchmarks, including 3D medical image segmentation. Hierarchical transformers (e.g., Swin Transformers) reintroduced several ConvNet priors and further enhanced the practical viability of adapting volumetric segmentation in 3D medical datasets. The effectiveness of hybrid approaches is largely credited to the large receptive field for non-local self-attention and the large number of model parameters. In this work, we propose a lightweight volumetric ConvNet, termed 3D UX-Net, which adapts the hierarchical transformer using ConvNet modules for robust volumetric segmentation. Specifically, we revisit volumetric depth-wise convolutions with large kernel size (e.g. starting from $7 imes7 imes7$) to enable the larger global receptive fields, inspired by Swin Transformer. We further substitute the multi-layer perceptron (MLP) in Swin Transformer blocks with pointwise depth convolutions and enhance model performances with fewer normalization and activation layers, thus reducing the number of model parameters. 3D UX-Net competes favorably with current SOTA transformers (e.g. SwinUNETR) using three challenging public datasets on volumetric brain and abdominal imaging: 1) MICCAI Challenge 2021 FLARE, 2) MICCAI Challenge 2021 FeTA, and 3) MICCAI Challenge 2022 AMOS. 3D UX-Net consistently outperforms SwinUNETR with improvement from 0.929 to 0.938 Dice (FLARE2021) and 0.867 to 0.874 Dice (Feta2021). We further evaluate the transfer learning capability of 3D UX-Net with AMOS2022 and demonstrates another improvement of $2.27\%$ Dice (from 0.880 to 0.900). The source code with our proposed model are available at https://github.com/MASILab/3DUX-Net.

研究の動機と目的

• パフォーマンスとモデルサイズのバランスを取る効率的な3Dセグメンテーションのバックボーンの必要性を動機づける。
• 大規模カーネル深さ方向畳み込みを用いて階層的なトランスフォーマー挙動を模倣する軽量な体積ConvNetを提案する。
• パラメータ数と正規化要件を削減しつつセグメンテーション精度を維持または向上させる。
• 監視学習および転移学習設定で公開脳および腹部データセットで強力な実証結果を示す。

提案手法

• 大きなカーネルサイズ（LK）を用いた体積深さ方向畳み込みを導入し、広い受容野を模倣する。
• TransformerブロックのMLPをポイントワイズ深さ方向畳み込みに置換し、パラメータを減らして特徴を広げる（DCS）。
• 深さ方向畳み込みを用いたインバーテッドボトルネック設計で、層間のチャンネル特徴を拡張・圧縮する。
• エンコーダブロックでバッチ正規化をレイヤー正規化に置換し、GELU活性化を用いる。
• 各段に2つのLKブロックを備えた4段階のエンコーダを構築し、スキップ接続を伴うConvNetベースのU字型デコーダを用いる。

実験結果

リサーチクエスチョン

• RQ1 大カーネル深さ方向畳み込みを備えた純粋なConvNetは、トランスフォーマーベースの3Dセグメンテーションの性能に匹敵するか、またはそれを上回るか。
• RQ2 提案された3D UX-Netは、より少ないパラメータ数と正規化層の削減で、同等かそれ以上の精度を達成するか。
• RQ3 公的ボリュームデータセットでの監督付き訓練と転移学習で、3D UX-Netはどのように性能を発揮するか。
• RQ4 カーネルサイズと深さ方向スケーリング（DCS）がデータセット間でのセグメンテーション性能に与える影響は。

主な発見

• 3D UX-Net は FeTA2021 で SwinUNETR を上回り（Dice 0.874 対 0.867）、FLARE2021 でも上回り（0.934 対 0.929 Dice）。
• AMOS2022 への転移学習で、3D UX-Net は Dice 0.900 を達成し、最良のトランスフォーマーベースのベースラインより 2.27% の改善。
• アブレーション研究は、カーネルサイズと深さ方向スケーリングが性能に影響を与え、LK サイズが 7x7x7 から 13x13x13 程度でデータセットごとに最適な利得を提供。
• 3D UX-Net は FeTA2021 で収束が速く、AMOS2022 で堅牢な転移学習挙動を示す。
• モデルは、いくつかのトランスフォーマーベースのベースラインと比較して、少ないパラメータ（53.0M）で競争力のある Dice スコアを達成。