[論文レビュー] U-Netmer: U-Net meets Transformer for medical image segmentation

Sheng He, Rina Bao+2 more|arXiv (Cornell University)|Apr 3, 2023

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

ひとこと要約

U-Netmerは、医用画像分類のためのU-NetとTransformerアーキテクチャの新規融合を提案し、1次元トークン化に伴う『トークンフラットネス』問題と、固定パッチサイズによる『スケール感受性』問題を、局所的パッチにU-Netを適用し、Transformerでグローバルなコンテキストをモデル化することで解決する。7つのデータセットおよびモダリティで最先端の性能を達成し、スケール間の分類差違を基に、教師なしでテスト画像の難易度順にランク付け可能な信頼性の高い信頼スコアを提供する。

ABSTRACT

The combination of the U-Net based deep learning models and Transformer is a new trend for medical image segmentation. U-Net can extract the detailed local semantic and texture information and Transformer can learn the long-rang dependencies among pixels in the input image. However, directly adapting the Transformer for segmentation has ``token-flatten" problem (flattens the local patches into 1D tokens which losses the interaction among pixels within local patches) and ``scale-sensitivity" problem (uses a fixed scale to split the input image into local patches). Compared to directly combining U-Net and Transformer, we propose a new global-local fashion combination of U-Net and Transformer, named U-Netmer, to solve the two problems. The proposed U-Netmer splits an input image into local patches. The global-context information among local patches is learnt by the self-attention mechanism in Transformer and U-Net segments each local patch instead of flattening into tokens to solve the `token-flatten" problem. The U-Netmer can segment the input image with different patch sizes with the identical structure and the same parameter. Thus, the U-Netmer can be trained with different patch sizes to solve the ``scale-sensitivity" problem. We conduct extensive experiments in 7 public datasets on 7 organs (brain, heart, breast, lung, polyp, pancreas and prostate) and 4 imaging modalities (MRI, CT, ultrasound, and endoscopy) to show that the proposed U-Netmer can be generally applied to improve accuracy of medical image segmentation. These experimental results show that U-Netmer provides state-of-the-art performance compared to baselines and other models. In addition, the discrepancy among the outputs of U-Netmer with different scales is linearly correlated to the segmentation accuracy which can be considered as a confidence score to rank test images by difficulty without ground-truth.

研究の動機と目的

• 視覚Transformerにおける『トークンフラットネス』問題に対処すること。これは、1次元トークン化の過程で局所的パッチ間の相互作用が失われるためである。
• パッチベースのTransformerにおける『スケール感受性』問題を克服すること。これは、固定パッチサイズのため、分類精度が顕著に変動することに起因する。
• 複数のパッチスケールをサポートする統一されたモデル構造を構築すること。これにより、共有パラメータによるより高いロバスト性が得られる。
• スケール間の分類一貫性に基づいて、テスト画像の信頼スコアを生成すること。これにより、教師なしで画像の難易度をランク付け可能となる。
• 多様な医用画像モダリティおよび臓器にわたる最先端の分類性能を実証すること。

提案手法

• 入力画像は複数のスケール（s = 1, 2, 4, 8）で重複のない局所的パッチに分割され、s=1は元の画像を意味する。
• 各局所的パッチはU-Netバックボーンを用いて独立に分類され、パッチ内での空間的相互作用が保持される。
• グローバルなコンテキストは、すべての局所的パッチの埋め込みに注目するTransformerエンコーダーによってモデル化される。
• 最終的な分類出力は、U-Netによるパッチ分類出力とグローバルアテンション特徴を統合することで得られる。
• 信頼スコアは、異なるスケール（例：s=1とs=2）からの分類マップ間のDice係数として計算され、分類精度の代理として機能する。
• モデルはスケール間で共有パラメータを用いてエンドツーエンドに訓練され、再トレーニングなしにマルチスケール推論が可能となる。

実験結果

リサーチクエスチョン

• RQ1 U-NetとTransformerのグローバル・ローカル統合は、標準的なU-Netやハイブリッドモデルに比べ、医用画像分類で優れた性能を発揮するか？
• RQ2 異なるスケール間の分類マップの差違は、実際の分類精度と相関を示すか？これにより、教師なしで信頼スコアを推定可能か？
• RQ3 提案手法は、再トレーニングなしに、さまざまなパッチスケールで一貫した性能を維持できるか？
• RQ4 スケール差違に基づく信頼スコアは、既存手法と比較して画像の難易度ランク付けにおいて優れているか？
• RQ5 このモデルは、多様な臓器および画像モダリティに一般化可能か？

主な発見

• U-Netmerは、7つの臓器と4つの画像モダリティ（MRI、CT、エラスソングラフィ、内視鏡）をカバーする7つの公的データセットで最先端の性能を達成した。
• 全評価データセットで4つのベースラインおよび6つのSOTAモデルを上回る一貫した改善を示した。
• 異なるスケール（例：s=1とs=2）間の分類マップの差違は、実際のDice精度と強い線形相関を示し、7つのデータセットのうち6つでピアソン相関係数が0.8を超えた。
• スケール差違に基づく信頼スコアは、6つのデータセットでProtoSegの平均分類能力スコアを上回り、画像の難易度ランク付けにおいて優れた性能を示した。
• 低信頼スコアの画像は、小規模でコントラストが低く、ぼやけたターゲット領域を示す傾向があり、高信頼スコアの画像は明確で大きく、明確に定義された構造を持つ。
• スケール間で一貫した予測を生成できる能力により、人間が関与するレビューのためのハードな分類ケースを信頼性高く同定できるようになった。