[論文レビュー] Medical Transformer: Gated Axial-Attention for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Poojan Oza+2 more|arXiv (Cornell University)|Feb 21, 2021

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

ひとこと要約

本稿では、医療画像セグメンテーションのためのゲート付き軸自己注意ネットワークであるMedical Transformer (MedT) を提案する。この手法は、軸自己注意に学習可能なゲート機構を統合することで、長距離特徴モデリングを向上させ、事前学習を必要とせずに3つの医療画像データセットで最先端の性能を達成する。特に、局所的・グローバルな学習戦略（LoGo）を用いることで、畳み込みモデルおよび既存のトランスフォーマーベースモデルを凌駕する。

ABSTRACT

Over the past decade, Deep Convolutional Neural Networks have been widely adopted for medical image segmentation and shown to achieve adequate performance. However, due to the inherent inductive biases present in the convolutional architectures, they lack understanding of long-range dependencies in the image. Recently proposed Transformer-based architectures that leverage self-attention mechanism encode long-range dependencies and learn representations that are highly expressive. This motivates us to explore Transformer-based solutions and study the feasibility of using Transformer-based network architectures for medical image segmentation tasks. Majority of existing Transformer-based network architectures proposed for vision applications require large-scale datasets to train properly. However, compared to the datasets for vision applications, for medical imaging the number of data samples is relatively low, making it difficult to efficiently train transformers for medical applications. To this end, we propose a Gated Axial-Attention model which extends the existing architectures by introducing an additional control mechanism in the self-attention module. Furthermore, to train the model effectively on medical images, we propose a Local-Global training strategy (LoGo) which further improves the performance. Specifically, we operate on the whole image and patches to learn global and local features, respectively. The proposed Medical Transformer (MedT) is evaluated on three different medical image segmentation datasets and it is shown that it achieves better performance than the convolutional and other related transformer-based architectures. Code: https://github.com/jeya-maria-jose/Medical-Transformer

研究の動機と目的

• 医療画像における長距離依存関係をモデル化する点で、畳み込みネットワークの限界を克服すること。
• 標準的なトランスフォーマーの効果的学習を妨げる医療画像分野のデータ不足の課題を克服すること。
• グローバルな文脈と局所的な詳細を明示的にモデル化することで、セグメンテーション精度を向上させるトランスフォーマーベースのアーキテクチャを開発すること。
• 全画像およびパッチレベルの表現を効果的に活用する学習戦略を設計し、小規模な医療画像データセットにおける性能を向上させること。

提案手法

• 2次元自己注意の注意重みに学習可能なゲートを追加するゲート付き軸自己注意メカニズムを提案し、関連する空間的文脈にのみ注目することで特徴表現を向上させる。
• 全画像（グローバルブランチ）と画像パッチ（ローカルブランチ）の両方を同時に学習する局所的・グローバル的（LoGo）学習戦略を導入し、グローバルな文脈と微細な局所特徴を両方学習可能にする。
• 軸自己注意を用いて2次元自己注意を2つの1次元操作に分解し、計算コストを低減しながらも空間モデリング能力を維持する。
• エンコーダーで標準的な畳み込み層の代わりに軸自己注意ブロックを搭載した、U-Netに類似したエンコーダー・デコーダー構造を採用する。
• 2段階の学習プロトコルを採用：最初の10エポックはゲートの更新を停止して特徴を安定化させた後、ベースの起動学習率0.001で全モジュールを同時に学習する。
• グローバルブランチを2つのエンコーダー／デコーダー層に制限し、ローカルブランチでパッチベース処理を実施することで、マルチブランチ構造にもかかわらず効率を維持する。

実験結果

リサーチクエスチョン

• RQ1 事前学習を必要とせずに、軸自己注意を備えたトランスフォーマーベースのアーキテクチャが、医療画像セグメンテーションにおいて長距離依存関係を効果的にモデル化できるか？
• RQ2 軸自己注意メカニズムに学習可能なゲートを導入することで、小規模な医療画像データセットにおける性能がどのように向上するか？
• RQ3 局所的・グローバル的（LoGo）学習戦略は、全画像またはパッチ単位での学習に比べて、特徴学習をどの程度向上させるか？
• RQ4 提案されたMedical Transformer（MedT）は、限られた医療データにおいて、既存の畳み込みモデルおよびトランスフォーマーベースモデルを上回るセグメンテーション精度を達成できるか？

主な発見

• Brain USデータセットでは、MedTがF1スコア88.84を達成し、U-Net（85.37）、Res-UNet（87.5）、Axial-UNet（88.39）を上回った。
• アブレーションスタディの結果、ゲート付き軸自己注意、LoGo学習、マルチブランチ構造の各要素が性能向上に寄与していることが確認された。
• グローバルエンコーダーをわずか2層に制限したLoGo戦略でも強力な性能を発揮しており、パッチレベルの学習と組み合わせることで、最小限のグローバルモデリングで十分であることが示された。
• パrameter数がわずか140万個（1.4M）であるにもかかわらず、元のU-Net（1250万個）やRes-UNet（532万個）よりも優れた性能を示しており、パrameter効率の高さが裏付けられた。
• 定性的な結果では、MedTはU-Net や Res-UNet が誤分類する複雑な背景テクスチャを有する超音波画像においても、背景領域を正しくセグメンテーションしていることが確認された。
• ImageNetの事前学習を必要とせず、TransUNet や TransFuse といった関連研究とは異なり、大規模な事前学習に依存しない状態で最先端の性能を達成した。