[論文レビュー] MEW-UNet: Multi-axis representation learning in frequency domain for medical image segmentation

Jiacheng Ruan, Mingye Xie+3 more|arXiv (Cornell University)|Oct 25, 2022

Advanced Neural Network Applications被引用数 15

ひとこと要約

MEW-UNet は、3つの空間軸に沿って2次元フーリエ変換を適用してグローバルな周波数ドメイン表現を学習する新規なマルチアックス外部重みブロックを提案する。これに深さ方向の畳み込みを組み合わせ、自己注意機構を置き換えるU-Netアーキテクチャを採用する。モデルは、SynapseデータセットにおいてHD95でMT-UNetを10.15mm上回り、最先端の性能を達成した。

ABSTRACT

Recently, Visual Transformer (ViT) has been widely used in various fields of computer vision due to applying self-attention mechanism in the spatial domain to modeling global knowledge. Especially in medical image segmentation (MIS), many works are devoted to combining ViT and CNN, and even some works directly utilize pure ViT-based models. However, recent works improved models in the aspect of spatial domain while ignoring the importance of frequency domain information. Therefore, we propose Multi-axis External Weights UNet (MEW-UNet) for MIS based on the U-shape architecture by replacing self-attention in ViT with our Multi-axis External Weights block. Specifically, our block performs a Fourier transform on the three axes of the input feature and assigns the external weight in the frequency domain, which is generated by our Weights Generator. Then, an inverse Fourier transform is performed to change the features back to the spatial domain. We evaluate our model on four datasets and achieve state-of-the-art performances. In particular, on the Synapse dataset, our method outperforms MT-UNet by 10.15mm in terms of HD95. Code is available at https://github.com/JCruan519/MEW-UNet.

研究の動機と目的

• 既存の医療画像セグメンテーションモデルが空間ドメイン表現にのみ焦点を当て、周波数ドメイン情報を見過ごしているという限界を是正すること。
• マルチアックスフーリエ変換を用いて周波数ドメイン特徴を統合することで、U-Netベースアーキテクチャにおけるグローバルなコンテキストモデリングを向上させること。
• 周波数ドメインの注目機構と深さ方向分離畳み込みを統合したブロックを用いて、グローバルおよびローカル特徴の学習をバランスさせること。
• ビジョントランスフォーマーにおける標準的な自己注意に依存せずに、周波数に敏感な注目機構を設計し、表現学習を強化すること。
• 周波数ドメイン不変性とマルチスケールコンテキストモデリングを統合することで、複数の医療画像ベンチマークで最先端のセグメンテーション性能を達成すること。

提案手法

• マルチアックス外部重み（MEW）ブロックは、チャネル次元に沿って入力特徴マップを4つのブランチに分割する。
• 3つのブランチは、高さ-幅、チャネル-幅、チャネル-高さの各軸に沿って2次元離散フーリエ変換（2D DFT）を適用し、特徴を周波数ドメインに変換する。
• 学習可能な外部重みを周波数ドメインに適用して周波数成分を調整し、グローバルコンテキストモデリングを強化する。
• 逆2次元フーリエ変換を適用して、重み付けされた周波数特徴を再び空間ドメインに戻す。
• 4番目のブランチは、ローカルな空間的詳細を保持するため、深さ方向分離畳み込みを適用する。
• 4つのブランチの出力を連結し、元の入力と残差接続することで、マルチアックス外部重みブロック（MEWB）を構築する。

実験結果

リサーチクエスチョン

• RQ1 マルチアックス周波数ドメイン表現の学習は、空間ドメインモデリングを越えて医療画像セグメンテーションを改善できるか？
• RQ2 複数の軸に沿って2次元フーリエ変換を用いた周波数ドメインでの学習は、モデルのグローバルコンテキストを捉える能力にどのように影響するか？
• RQ3 統合されたブロック内で周波数ドメイン注目機構とローカル畳み込みを組み合わせた際のセグメンテーションへの寄与度は何か？
• RQ4 提案されたMEW-UNetは、医療画像セグメンテーションタスクにおいて、従来のハイブリッドCNN-ViTおよび純粋なViTモデルを上回る性能を示すか？
• RQ5 周波数ドメインに学習可能な外部重みを組み込むことで、モデルの一般化性能と性能にどのような影響が生じるか？

主な発見

• MEW-UNet は、ISIC17、ISIC18、Synapse、ACDC の4つの公的医療画像セグメンテーションデータセットで最先端の性能を達成した。
• Synapseデータセットでは、MT-UNetに比べてHD95が10.15mm向上し、DSCが0.33%向上した。境界局在の優位性が明確に示された。
• ACDCデータセットでは、MEW-UNetがDSCスコア91.00%を達成し、MT-UNet（90.43%）とSwin-Unet（88.07%）を上回った。
• アブレーションスタディの結果、3つの軸に沿った周波数ドメイン演算を追加するごとにmIoUとDSCが順次向上し、完全なMEWBではISIC18でmIoUが81.90%、DSCが90.05%を達成した。
• アブレーションスタディでGroupNormをBatchNormに置き換えた場合、性能が低下した。これは、このアーキテクチャにおいてGroupNormがより効果的であることを示している。
• Synapseデータセットの可視化結果から、MEW-UNetは特に臓器の境界周辺で、より鋭く正確なセグメンテーションマスクを生成していることが確認された。