[論文レビュー] UNet-2022: Exploring Dynamics in Non-isomorphic Architecture

Jiansen Guo, Hong-Yu Zhou+2 more|arXiv (Cornell University)|Oct 27, 2022

Radiomics and Machine Learning in Medical Imaging被引用数 8

ひとこと要約

UNet-2022 は、自己注意と畳み込みモジュールを並列処理することで、動的空間的およびチャネル別特徴重み付けを共同で活用する、非同型なU-Netアーキテクチャを提案する。このアーキテクチャにより、臓器多臓器および心臓画像セグメンテーションを含む、複数の医療画像セグメンテーションタスクで最先端の性能を達成し、DiceスコアでnnUNetを最大4%上回った。

ABSTRACT

Recent medical image segmentation models are mostly hybrid, which integrate self-attention and convolution layers into the non-isomorphic architecture. However, one potential drawback of these approaches is that they failed to provide an intuitive explanation of why this hybrid combination manner is beneficial, making it difficult for subsequent work to make improvements on top of them. To address this issue, we first analyze the differences between the weight allocation mechanisms of the self-attention and convolution. Based on this analysis, we propose to construct a parallel non-isomorphic block that takes the advantages of self-attention and convolution with simple parallelization. We name the resulting U-shape segmentation model as UNet-2022. In experiments, UNet-2022 obviously outperforms its counterparts in a range segmentation tasks, including abdominal multi-organ segmentation, automatic cardiac diagnosis, neural structures segmentation, and skin lesion segmentation, sometimes surpassing the best performing baseline by 4%. Specifically, UNet-2022 surpasses nnUNet, the most recognized segmentation model at present, by large margins. These phenomena indicate the potential of UNet-2022 to become the model of choice for medical image segmentation.

研究の動機と目的

• 自己注意と畳み込みの相補的利点が医療画像セグメンテーションにおいてどのように機能するかを直感的に説明すること。
• 自己注意と畳み込みを組み合わせる既存のハイブリッドアーキテクチャに、明確な根拠なしに統合されているという解釈不能性の問題を解決すること。
• 空間的およびチャネル次元の両方で動的重み割り当てを同時に可能にする新しいブロックを設計すること。
• シンプルでありながら効果的なアーキテクチャ的イノベーションを通じて、多様な医療画像処理タスクにおけるセグメンテーション性能を向上させること。
• 提案されたUNet-2022フレームワークにおけるImageNetベースの事前学習と推論戦略チューニングの有効性を検証すること。

提案手法

• 自己注意と畳み込み（DwConv）モジュールを独立して処理する並列非同型ブロック（PIブロック）を提案する。
• 空間的およびチャネル別に適応的な重みを割り当てる動的重み割り当てメカニズムを導入し、両モジュールの利点を統合する。
• パrameter数を削減し、過学習を軽減しながらも性能を維持するために、深度方向畳み込みを活用する。
• スキップ接続を備えたU-Net風のエンコーダデコーダ構造を採用し、標準的なエンコーダを提案されたPIブロックに置き換える。
• 一般化性能を向上させるために、ImageNetベースの事前学習を適用する。特にデータ量が少ない状況での有効性が顕著である。
• 特徴量の集約を向上させるために、小さなステップサイズ（0.2× クロップサイズ）を用いたスライディングウィンドウ推論を最適化する。

実験結果

リサーチクエスチョン

• RQ1 自己注意と畳み込みをハイブリッドに組み合わせたアーキテクチャが、純粋な畳み込みモデルや自己注意ベースのモデルよりも医療画像セグメンテーションで優れた性能を発揮する理由は何か？
• RQ2 自己注意と畳み込みの動的重み割り当てメカニズムを正式に比較し、統一されたフレームワークで活用する方法は何か？
• RQ3 自己注意とDwConvの並列的かつ非同型な処理を導入することで、多様な医療画像処理タスクにおけるセグメンテーション性能にどのような影響を与えるか？
• RQ4 ImageNetベースの事前学習は、学習から開始する場合と比較して、提案されたUNet-2022モデルの性能をどの程度向上させるか？
• RQ5 スライディングウィンドウ推論における推論ステップサイズの変更が、最終的なセグメンテーション品質に与える影響は何か？

主な発見

• 提案されたPIブロックは、Synapseデータセットにおいて平均Diceスコア86.46、HD95 11.34を達成し、Swin Transformerブロック（84.42 DSC、20.74 HD95）およびConvNeXtブロック（84.96 DSC、16.60 HD95）を上回った。
• PIブロックから自己注意モジュールを削除すると、平均DSCは85.20に低下し、HD95は14.96に上昇する。これは空間的ダイナミクスの重要性を裏付けている。
• ImageNetベースの事前学習により、UNet-2022の平均DSCは1.5ポイント上昇し、HD95は約5mm改善された。学習から開始する場合と比較して顕著な向上が確認された。
• 推論ステップサイズを0.2× クロップサイズにすることで、HD95は0.5× クロップサイズと比較して2.6mm改善された。これは、より密な推論サンプリングの利点を示している。
• 腹部多臓器および皮膚病変セグメンテーションを含む、複数のタスクにおいて、UNet-2022は広く採用されているnnUNetモデルを最大4%のDiceスコアで上回った。
• アブレーションスタディの結果、自己注意とDwConvの並列統合は、単独で使用する場合や逐次的に使用する場合よりも効果的であることが確認された。