[論文レビュー] Narrowing the semantic gaps in U-Net with learnable skip connections: The case of medical image segmentation

Haonan Wang, Peng Cao+3 more|arXiv (Cornell University)|Dec 23, 2023

COVID-19 diagnosis using AI被引用数 4

ひとこと要約

本稿では、U-Netベースの医療画像セグメンテーションフレームワークUDTransNetを提案する。この手法は、標準的なスキップ接続の代わりに学習可能なアテンション機構を導入し、エンコーダーとデコーダーの特徴量間の意味的ギャップを縮小する。マルチスケール特徴量統合のための二重アテンション変換器（DAT）と、クロスステージの一貫性を向上させるデコーダー誘導型再キャリブレーションアテンション（DRA）モジュールを導入することで、パrameter数を減らしながらも、多数の医療画像ベンチマークで最先端の性能を達成した。

ABSTRACT

Most state-of-the-art methods for medical image segmentation adopt the encoder-decoder architecture. However, this U-shaped framework still has limitations in capturing the non-local multi-scale information with a simple skip connection. To solve the problem, we firstly explore the potential weakness of skip connections in U-Net on multiple segmentation tasks, and find that i) not all skip connections are useful, each skip connection has different contribution; ii) the optimal combinations of skip connections are different, relying on the specific datasets. Based on our findings, we propose a new segmentation framework, named UDTransNet, to solve three semantic gaps in U-Net. Specifically, we propose a Dual Attention Transformer (DAT) module for capturing the channel- and spatial-wise relationships to better fuse the encoder features, and a Decoder-guided Recalibration Attention (DRA) module for effectively connecting the DAT tokens and the decoder features to eliminate the inconsistency. Hence, both modules establish a learnable connection to solve the semantic gaps between the encoder and the decoder, which leads to a high-performance segmentation model for medical images. Comprehensive experimental results indicate that our UDTransNet produces higher evaluation scores and finer segmentation results with relatively fewer parameters over the state-of-the-art segmentation methods on different public datasets. Code: https://github.com/McGregorWwww/UDTransNet.

研究の動機と目的

• U-Netにおける固定スキップ接続の限界を調査すること。
• すべてのスキップ接続が同じ程度に有用であるわけではないこと、最適な組み合わせはデータセットによって異なることの特定。
• マルチスケールエンコーダー特徴量間およびエンコーダー・デコーダーステージ間の意味的ギャップを解消すること。
• U-Netアーキテクチャ全体にわたる特徴量の適合性を向上させる、学習可能なアテンションベースのメカニズムの設計。
• パrameter数を削減しつつ、最先端のセグメンテーション性能を達成すること。

提案手法

• マルチスケールエンコーダー特徴量にわたるチャネルおよび空間的依存性をモデル化するため、チャネル別統合アテンション（CFA）と空間別選択アテンション（SSA）を組み合わせた二重アテンション変換器（DAT）モジュールを提案。
• CFAでは自己アテンションを用いて異なるエンコーダーステージ間のチャネルをアライメントし、SSAではクロスアテンションを用いてスケール間で重要な空間領域を強調。
• デコーダー特徴量を用いてDAT出力を再キャリブレーションすることで、意味的不一致を低減するデコーダー誘導型再キャリブレーションアテンション（DRA）モジュールを導入。
• U-Netボトムアップアーキテクチャのスキップ接続にDATとDRAを統合し、アテンションベースの特徴量統合を可能にするエンドツーエンド学習を実現。
• 主エンコーダーやデコーダーにアテンションを適用するのではなく、CNNベースのエンコーダーとスキップ接続におけるトランスフォーマーベースのアテンションモジュールをハイブリッドに統合。
• アテンションヘッド数やトランスフォーマーレイヤー数などのハイパーパramータを最適化し、4ヘッドおよび4レイヤーで最適な設定が得られることを特定。

実験結果

リサーチクエスチョン

• RQ1 U-Netにおけるすべてのスキップ接続が、異なる医療画像セグメンテーションタスクにおいて同等に有効であるか？
• RQ2 マルチスケールエンコーダー特徴量間およびエンコーダー・デコーダーステージ間の意味的ギャップが、セグメンテーション性能に与える影響は何か？
• RQ3 スキップ接続に学習可能なアテンション機構を適用することで、特徴量のアライメントが向上し、セグメンテーション精度が向上するか？
• RQ4 アテンションをスキップ接続に統合することは、エンコーダーまたはデコーダーに適用するのと比較して、より効果的か？
• RQ5 提案されたアテンションモジュールにおけるアテンションヘッド数およびレイヤー数の最適な設定は何か？

主な発見

• UDTransNetは3つの公開データセットで最先端のDiceスコアを達成した：GlaSでは91.03±0.56、Synapseでは89.28±0.58を記録し、既存手法を上回った。
• 同じバックボーン（ResNet-34）を用いた場合、UDTransNetはGlaSで91.03±0.56のDiceスコアを達成し、MCTrans（88.44±0.74）およびTransUNet（87.41±1.15）を顕著に上回った。
• Swin Transformerをバックボーンとして用いた場合でも、UDTransNetはGlaSで88.03±1.27のDiceスコアを達成し、同じバックボーンを用いたSwin-UNet（89.04±0.88）を上回った。
• ハイパーパramータの変更に対してもモデルは頑健であり、最適性能は4ヘッドおよび4レイヤーで得られ、これらの設定に対して低い感度を示した。
• アブレーションスタディにより、DATおよびDRAモジュールが両方とも不可欠であることが確認され、特にDRAがエンコーダーとデコーダー間の特徴量の一貫性を顕著に向上させた。
• 最先端のモデルと比較して、パrameter数を削減しながらも高い性能を達成したため、効率性と有効性の両面で優れた性能を示した。