[論文レビュー] Large-kernel Attention for Efficient and Robust Brain Lesion Segmentation

Liam Chalcroft, Ruben Lourenço Pereira+6 more|arXiv (Cornell University)|Aug 14, 2023

Advanced Neural Network Applications被引用数 4

ひとこと要約

本稿では、3D脳病変セグメンテーションのためのU-Netを向上させるために、トランスフォーマーのグローバルな受容 field とCNNのパラメータ効率性および局所パターンへのインダクティブバイアスを組み合わせたハイブリッド畳み込み-アテンションモジュールであるLarge-kernel Attention (LKA) を提案する。LKAはドメインシフトに対して優れたロバスト性を示し、BraTSおよびATLASデータセットでも競争力ある性能を達成しており、特に形状ベースの一般化および分布外一般化において優れた性能を発揮する。

ABSTRACT

Vision transformers are effective deep learning models for vision tasks, including medical image segmentation. However, they lack efficiency and translational invariance, unlike convolutional neural networks (CNNs). To model long-range interactions in 3D brain lesion segmentation, we propose an all-convolutional transformer block variant of the U-Net architecture. We demonstrate that our model provides the greatest compromise in three factors: performance competitive with the state-of-the-art; parameter efficiency of a CNN; and the favourable inductive biases of a transformer. Our public implementation is available at https://github.com/liamchalcroft/MDUNet .

研究の動機と目的

• 脳病変セグメンテーションにおいて、長距離の空間的依存関係を捉えることが不可欠であるが、標準的なCNNがそのような依存関係を十分に捉えられていないという限界を解消すること。
• 医用画像における標準的なビジョントランスフォーマーの計算非効率性および局所的なインダクティブバイアスの欠如を克服すること。
• MRIコントラストシーケンスにおけるドメインシフトに強い、パフォーマンス、パラメータ効率性、ロバスト性のバランスを取ったモデルを開発すること。
• 大カーネルアテンションを有するハイブリッドアーキテクチャが、トランスフォーマーのインダクティブバイアスを実現しながらも、CNNと同等の効率性を維持できるかどうかを調査すること。

提案手法

• 大カーネルアテンション（LKA）に基づく新規な全畳み込み型トランスフォーマーブロックを提案し、大カーネルを深さ方向畳み込み、拡張深さ方向畳み込み、ポイントワイド畳み込みに分解する。
• 標準的な自己アテンションまたは畳み込みブロックに置き換えることで、LKAモジュールをU-Netアーキテクチャに統合し、パラメータ数を削減しつつグローバルな文脈モデリングを可能にする。
• 要因分解されたアテンションメカニズムを採用：X₁ = dDWConv(DWConv(X))、X₂ = PConv(X₁) により、空間的およびチャネル混合を効率的に行う。
• ガウスノイズと強度のジャイタリングを含むデータオーグメンテーションを用いて、マルチモodal MRIデータ上でエンドツーエンドでモデルを学習する。損失関数には交差エントロピーとDice損失を併用する。
• BraTSでの推論安定性と速度向上のため、テスト時オーグメンテーションと半精度推論を採用する。
• ISLES 2015の未学習なMRIシーケンス（T1、T2、DWI、FLAIR）へのゼロショット転送を用いて一般化性能を評価する。

(a) LKA

実験結果

リサーチクエスチョン

• RQ1 大カーネルアテンション機構は、計算効率性を維持しつつ、3D脳病変セグメンテーションにおける長距離依存関係を効果的にモデル化できるか？

主な発見

• BraTS 2021の公的テストセットにおいて、LKA-Eモデルは全腫瘍（WT）のDiceスコアが0.939を達成し、CNNベースのnnUNet（0.902）およびSwinベースのモデル（0.938）を上回った。HD95においても統計的に有意な改善（2.449 vs. 3.606、p < 0.0001）を示した。
• ATLASデータセットでは、LKAモデルが全指標で競争力ある性能を維持し、強化腫瘍（ET）のDiceスコアが0.888であった。平均表面距離（AVD）においても、CNNベースラインを顕著に上回った。
• ISLES 2015における分布外（OOD）評価では、LKAモデルはMRIコントラストのシフトに対して優れたロバスト性を示し、T1、T2、DWI、FLAIRシーケンスすべてで高い性能を維持した。一方、nnUNetベースラインは著しく性能を低下させた。
• ガウスノイズによる画像ぼかし（σ = 0.5から2.0）に対して、LKAモデルは安定したDiceおよびHD95スコアを示し、Swinトランスフォーマーに近い形状バイアスを示した。これに対して、nnUNet（テクスチャバイアス）は性能が著しく低下した。
• 腫瘍コア（TC）の中央値HD95は2.000を達成し、Swinベースのモデル（2.236）およびCNN（2.236）を顕著に上回った。これは表面精度の向上を示している。
• アブレーションスタディにより、LKAにおけるカーネルサイズと拡張率を変更することで、インダクティブバイアスを微調整でき、タスク固有の最適化が可能であることが確認された。

(b) Attention