[論文レビュー] Learning Multi-axis Representation in Frequency Domain for Medical Image Segmentation

Jiacheng Ruan, Jingsheng Gao+2 more|arXiv (Cornell University)|Dec 28, 2023

Advanced Neural Network Applications被引用数 4

ひとこと要約

本稿では、3つの空間的・チャネル軸に沿って2次元離散フーリエ変換を適用し、周波数ドメインにおける多軸表現を学習するための新規なMulti-axis External Weightsブロックを用いて、医療画像セグメンテーションを向上させるU-Netベースのアーキテクチャ、MEW-UNetを提案する。周波数ドメインでの学習可能重みによる特徴の調制により、より包括的なグローバルおよびローカルなコンテキストを捉えることができ、Synapse, ACDC, ISIC17, ISIC18の4つの医療画像データセットで最先端の性能を達成し、ISIC18では平均Diceスコア91.00%を記録した。

ABSTRACT

Recently, Visual Transformer (ViT) has been extensively used in medical image segmentation (MIS) due to applying self-attention mechanism in the spatial domain to modeling global knowledge. However, many studies have focused on improving models in the spatial domain while neglecting the importance of frequency domain information. Therefore, we propose Multi-axis External Weights UNet (MEW-UNet) 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 features and assigns the external weight in the frequency domain, which is generated by our External 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, including Synapse, ACDC, ISIC17 and ISIC18 datasets, and our approach demonstrates competitive performance, owing to its effective utilization of frequency domain information.

研究の動機と目的

• 既存のビジョントランスフォーマーに基づくモデルが空間ドメインのモデリングに特化しており、周波数ドメインの情報を無視するという限界を是正すること。
• 複数の軸にわたる周波数ドメイン信号を活用することで、医療画像セグメンテーションにおけるグローバルコンテキストモデリングを向上させること。
• 周波数ドメイン表現とローカル空間的特徴を効果的に統合する学習可能なメカニズムを開発すること。
• 多軸周波数ドメイン解析が、単一軸解析に比べて、とくに微細な病理的構造に対してより特徴的な信号パターンを提供することを実証すること。

提案手法

• Multi-axis External Weightsブロック（MEWB）は、高さ-幅、チャネル-幅、チャネル-高さの3軸に沿って、それぞれ独立に2次元離散フーリエ変換（2D DFT）を適用し、多軸周波数表現を抽出する。
• 外部重み生成器を介して生成された学習可能な外部重みが、周波数ドメイン特徴に適用され、グローバルコンテキスト情報を調制する。
• 4番目のブランチでは、深さ方向可分畳み込み（DW畳み込み）を適用してローカル空間的詳細を保持し、ローカルおよびグローバル表現の補完的学習を確保する。
• MEWBは、U-Netのエンコーダ-デコーダアーキテクチャ内におけるビジョントランスフォーマーの自己注意モジュールに置き換えられ、MEW-UNetモデルが構築される。
• 重み付き調制後の周波数ドメイン特徴マップは、逆DFTを用いて空間ドメインに変換され、エンドツーエンド学習が可能になる。
• 標準的な指標（Diceスコア、mIoUなど）を用いて、4つの公開医療画像セグメンテーションデータセットでエンドツーエンド学習が実施される。

実験結果

リサーチクエスチョン

• RQ1 単一軸または空間のみのアプローチに比べて、多軸周波数ドメイン表現の学習が医療画像セグメンテーションモデルの性能向上に寄与するか。
• RQ2 周波数ドメインに学習可能な外部重みを統合することで、モデルの長距離依存性およびコンテキスト情報の捉え方にはどのような影響を与えるか。
• RQ3 セグメンテーションタスクにおける周波数ドメイン特徴とローカル空間的特徴の相対的寄与度は何か。また、それらはどのように相互作用するか。
• RQ4 提案されたMEW-UNetは、多様なデータセットにおいて、既存の最先端モデルを上回る性能を示すか。

主な発見

• MEW-UNetはISIC18データセットで平均Diceスコア91.00%を達成し、MT-UNet（90.43%）やSwin-Unet（88.07%）といった先行最先端モデルを上回った。
• アブレーションスタディの結果、DW畳み込みブランチを削除するとmIoUが0.61%低下し、ローカル特徴学習の重要性が裏付けられた。
• 外部重み生成器における内部表現ブロック（IRB）を除いたランダム重みの使用では、mIoUが1.31%低下し、学習可能な重み生成メカニズムの必要性が示された。
• 3つの周波数軸とDWブランチを備えた完全なMEWBは、ISIC18で81.90%のmIoUを達成し、単一軸またはDWなしのベースラインに比べ顕著に優れた性能を示した。
• 可視化結果から、MEW-UNetはMT-UNet や TransFuseに比べ、特に小さなまたは不規則な病変を有する困難なケースにおいて、より正確で完全な境界を生成することが分かった。
• 周波数ドメイン信号強度解析の結果、異なる組織タイプにおいて非重複な信号パターンが多軸DFTによって明らかになったことから、より優れた判別性を実現する設計選択の妥当性が裏付けられた。