[論文レビュー] SegStitch: Multidimensional Transformer for Robust and Efficient Medical Imaging Segmentation

Shengbo Tan, Zeyu Zhang+6 more|arXiv (Cornell University)|Aug 1, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用数 3

ひとこと要約

SegStitchは、軸方向パッチベースの自己注意機構とノイズ除去常微分方程式（ODE）ブロックを組み合わせることで、長距離特徴モデリングと耐性性を向上させる、3D医療画像セグメンテーションの新規フレームワークを提案する。UNETRと比較して、パrameter数を36.7%削減し、FLOPsを10.7%削減しながら、BTCVではmDSCが11.48%向上、ACDCでは6.71%向上を達成し、最先端の性能を実現した。

ABSTRACT

Medical imaging segmentation plays a significant role in the automatic recognition and analysis of lesions. State-of-the-art methods, particularly those utilizing transformers, have been prominently adopted in 3D semantic segmentation due to their superior performance in scalability and generalizability. However, plain vision transformers encounter challenges due to their neglect of local features and their high computational complexity. To address these challenges, we introduce three key contributions: Firstly, we proposed SegStitch, an innovative architecture that integrates transformers with denoising ODE blocks. Instead of taking whole 3D volumes as inputs, we adapt axial patches and customize patch-wise queries to ensure semantic consistency. Additionally, we conducted extensive experiments on the BTCV and ACDC datasets, achieving improvements up to 11.48% and 6.71% respectively in mDSC, compared to state-of-the-art methods. Lastly, our proposed method demonstrates outstanding efficiency, reducing the number of parameters by 36.7% and the number of FLOPS by 10.7% compared to UNETR. This advancement holds promising potential for adapting our method to real-world clinical practice. The code will be available at https://github.com/goblin327/SegStitch

研究の動機と目的

• 標準的なビジョントランスフォーマーが3D医療画像セグメンテーションにおいて抱える課題、特に高い計算コストと劣った局所特徴モデリング性能を是正すること。
• ハイブリッド自己注意機構を用いて、3D臓器および病変セグメンテーションにおける長距離依存関係の学習と意味的整合性を向上させること。
• 神経ODEブロックの統合により、連続的な特徴マップの最適化を可能とし、モデルの耐性性と一般化性能を向上させること。
• セグメンテーションの正確性を損なわず、モデルの複雑さと推論コストを低減することで、臨床応用を可能とすること。

提案手法

• SegStitchは、3Dボリューム全体の注目を避けるために、軸方向パッチ単位のクエリを用いることで、空間的および意味的整合性を維持し、計算量を削減する。
• 粗大および細粒度の自己注意を組み合わせた二重パスアテンション機構を導入し、グローバルな文脈と局所的詳細の両方をよりよく捉える。
• 推論中に連続的かつ動的に行われる特徴マップの更新を可能とする、ノイズ除去ODEブロック（nmODEに基づく）をモデルに統合する。
• ODEブロックは、特徴の改善に寄与する学習されたベクトル場（例：sin²、sin、tanh）を用いる。これにより、入力の摂動やデータノイズに対する耐性が向上する。
• パッチ間で共通のクエリを用いることで、パッチ間の意味的整合性が向上し、セグメンテーションの一貫性が向上する。
• アブレーションスタディをODEの変種について実施しながら、交差エントロピー損失とDice損失を用いて、3D医療画像データセット上でエンドツーエンドで訓練する。

実験結果

リサーチクエスチョン

• RQ1 共通クエリを用いた軸方向パッチベースのアテンションは、3D医療画像セグメンテーションにおける意味的整合性の向上と計算コストの低減に寄与するか？
• RQ2 ODEブロックの統合は、ノイズや摂動を含む医療スキャンに対して、モデルの耐性性と特徴の最適化をどのように向上させるか？
• RQ3 提案されたSegStitchフレームワークは、UNETRのような最先端のトランスフォーマーと比較して、セグメンテーションの正確性と効率性の面でどの程度優れているか？
• RQ4 sin²、シグモイド、tanhなどのODEの式表現のうち、どの形式が訓練の安定性、一般化性能、セグメンテーション性能のバランスを最も良く実現するか？

主な発見

• BTCVデータセットでは、最先端手法と比較して、平均Dice類似係数（mDSC）が11.48%向上した。
• ACDCデータセットでは、既存のSOTAアプローチと比較してmDSCが6.71%向上し、挑戦的な心臓セグメンテーションにおいて優れた性能を示した。
• UNETRと比較して、パrameter数を36.7%削減し、FLOPsを10.7%削減した。これは顕著な効率性の向上を示している。
• ODEブロックにsin²活性化関数を用いた場合、ACDCデータセットで最良の検証損失（-0.8895）とmDSC（93.32%）を達成し、シグモイドおよびオートエンコーダーのアブレーションを上回った。
• 定性的な結果から、右心室や腎臓のような細長い臓器の輪郭予測が顕著に向上しており、長距離依存関係モデリングの向上が確認された。
• アテンション重みの可視化により、モデルが顕著な細長い構造に注目していることが確認され、有効なグローバル特徴学習が行われていることが示された。