[論文レビュー] Medical Image Segmentation Using Squeeze-and-Expansion Transformers

Shaohua Li, Xiuchao Sui+4 more|arXiv (Cornell University)|May 20, 2021

Advanced Neural Network Applications参考文献 43被引用数 23

ひとこと要約

本論文は、スケーリング・エクスパンション・トランスフォーマーを用いて広い有効受容 field を実現しながらも、高い空間分解能を維持するトランスフォーマー基盤の医療画像セグメンテーションフレームワーク、Segtranを提案する。学習可能な正弦波位置エンコーディングとデュアルパス・アテンション機構を統合することで、Segtranは2Dおよび3Dの医療画像処理タスクにおいてU-Netおよびその変種を上回り、最小限のハイパーパrameterチューニングで最先端の精度と強力なクロスドメイン一般化性能を示した。

ABSTRACT

Medical image segmentation is important for computer-aided diagnosis. Good segmentation demands the model to see the big picture and fine details simultaneously, i.e., to learn image features that incorporate large context while keep high spatial resolutions. To approach this goal, the most widely used methods -- U-Net and variants, extract and fuse multi-scale features. However, the fused features still have small "effective receptive fields" with a focus on local image cues, limiting their performance. In this work, we propose Segtran, an alternative segmentation framework based on transformers, which have unlimited "effective receptive fields" even at high feature resolutions. The core of Segtran is a novel Squeeze-and-Expansion transformer: a squeezed attention block regularizes the self attention of transformers, and an expansion block learns diversified representations. Additionally, we propose a new positional encoding scheme for transformers, imposing a continuity inductive bias for images. Experiments were performed on 2D and 3D medical image segmentation tasks: optic disc/cup segmentation in fundus images (REFUGE'20 challenge), polyp segmentation in colonoscopy images, and brain tumor segmentation in MRI scans (BraTS'19 challenge). Compared with representative existing methods, Segtran consistently achieved the highest segmentation accuracy, and exhibited good cross-domain generalization capabilities. The source code of Segtran is released at https://github.com/askerlee/segtran.

研究の動機と目的

• U-Netおよびその変種が、小さな有効受容 field により長距離のコンテキストを捉えきれないという限界に対処すること。
• トランスフォーマーの無限大の有効受容 field を活用しながらも、高い空間分解能を維持することで医療画像セグメンテーションを向上させること。
• 自然言語事前学習からの制限を克服するため、医療画像セグメンテーションに特化したトランスフォーマー・アーキテクチャを設計すること。
• 新しい位置エンコーディング方式による連続性のインダクティブバイアスを組み込むことで、クロスドメイン一般化を強化すること。
• 多様な医療画像処理タスクにおいて、最小限のアーキテクチャ的変更とハイパーパrameterチューニングで最先端のパフォーマンスを達成すること。

提案手法

• 計算負荷の軽減と訓練の安定化を図るため、自己注意行列を正則化する「スケーティング・アテンション・ブロック」を備えたスケーリング・エクスパンション・トランスフォーマーを提案する。
• 特徴を複数のヘッドに投影し、アテンション後に統合することで多様な表現を学習する「エクスパンション・ブロック」を導入する。
• 標準的な学習可能または固定正弦波エンコーディングと比較して、画像データに対するインダクティブバイアスを向上させるために、空間的連続性を強制する学習可能な正弦波位置エンコーディングを採用する。
• 空間分解能の維持と回復を図るため、トランスフォーマーブロックの前後にも特徴ピラミッドネットワーク（FPN）を統合したU-Netに類似したエンコーダー・デコーダー構造を採用する。
• 初期特徴を抽出するために、CNNバックボーン（例：ResNet-101 または EfficientNet-B4）を採用し、その特徴をシーケンスにフラット化してトランスフォーマー処理に供する。
• ピクセル座標に基づく位置エンコーディングを適用することで、自己注意計算中に空間構造を保持する。

実験結果

リサーチクエスチョン

• RQ1 強化されたアテンション機構を備えたトランスフォーマー基盤アーキテクチャは、コンテキスト捕捉が制限される医療画像セグメンテーションタスクにおいて、U-Netおよびその変種を上回ることができるか？
• RQ2 スケーリング・エクスパンション・トランスフォーマー設計は、医療画像における標準的なトランスフォーマーと比較して、特徴表現とセグメンテーション精度を向上させるか？
• RQ3 提案された学習可能な正弦波位置エンコーディングは、医療画像セグメンテーションにおけるインダクティブバイアスとモデル一般化性能をどの程度向上させるか？
• RQ4 REFUGE20 と RIM-One のような、異なる画像特性を持つデータセットに対して、Segtran はどの程度のクロスドメイン一般化性能を示すか？
• RQ5 ImageNet での事前学習は、医療画像セグメンテーションモデルにどのような影響を及ぼし、異なるアーキテクチャにおいてパフォーマンスにどのように寄与するか？

主な発見

• Segtran は、REFUGE’20 チャレンジ（網膜神経線維層／杯領域セグメンテーション）で最高のDiceスコアを記録し、テストセットにおける平均Diceは 0.938 に達し、上位5チームの一つとなった。
• BraTS’19 チャレンジ（3D脳腫瘍セグメンテーション）では、U-Net や DeepLabV3+ を上回り、バリデーションセットにおける平均Diceスコアは 0.875 を達成した。
• EfficientNet-B4 をバックボーンとして使用した場合、Segtran は全タスクの平均で 0.909 のDiceスコアを記録し、U-Net（0.875）および TransU-Net（0.901）を顕著に上回った。
• Segtran は最も優れたクロスドメイン一般化性能を示し、REFUGE20 から RIM-One へのドメイン転送において、Diceスコアの低下が14.2%にとどまった。これに対して U-Net は18.2%、DeepLabV3+ は19.3%の低下を示した。
• ImageNet での事前学習により、Segtran の平均Diceスコアが約2.5%向上し、低データ環境下での事前学習の強力な利点を示した。
• ResNet-101 を使用した場合、FLOPs やパラメータ数が増加したが、EfficientNet-B4 と組み合わせることで、FLOPs を 71.3G に、パラメータ数を 93.1M に削減し、最適な効率性とパフォーマンスを達成した。