[論文レビュー] AgileFormer: Spatially Agile Transformer UNet for Medical Image Segmentation

Peijie Qiu, Yang Jin+3 more|arXiv (Cornell University)|Mar 29, 2024

COVID-19 diagnosis using AI被引用数 11

ひとこと要約

AgileFormer は ViT-UNet に変形可能パッチ埋め込み、時空間的に動的な自己注意、そして多層級の変形可能位置エンコーディングを導入し、3 つのデータセットにわたる医用画像分割の最先端性能と高いスケーラビリティを実現します。

ABSTRACT

In the past decades, deep neural networks, particularly convolutional neural networks, have achieved state-of-the-art performance in a variety of medical image segmentation tasks. Recently, the introduction of the vision transformer (ViT) has significantly altered the landscape of deep segmentation models. There has been a growing focus on ViTs, driven by their excellent performance and scalability. However, we argue that the current design of the vision transformer-based UNet (ViT-UNet) segmentation models may not effectively handle the heterogeneous appearance (e.g., varying shapes and sizes) of objects of interest in medical image segmentation tasks. To tackle this challenge, we present a structured approach to introduce spatially dynamic components to the ViT-UNet. This adaptation enables the model to effectively capture features of target objects with diverse appearances. This is achieved by three main components: extbf{(i)} deformable patch embedding; extbf{(ii)} spatially dynamic multi-head attention; extbf{(iii)} deformable positional encoding. These components were integrated into a novel architecture, termed AgileFormer. AgileFormer is a spatially agile ViT-UNet designed for medical image segmentation. Experiments in three segmentation tasks using publicly available datasets demonstrated the effectiveness of the proposed method. The code is available at \href{https://github.com/sotiraslab/AgileFormer}{https://github.com/sotiraslab/AgileFormer}.

研究の動機と目的

• セグメンテーションタスクで形状やサイズが異なる医用画像ターゲットのばらつきに対処する。
• 局所化とグローバルコンテキストの処理を改善する時空間的に動的な ViT-UNet の開発。
• 表現学習を強化するための変形可能パッチ埋め込みと変形可能位置エンコーディングを提案する。

提案手法

• 剛性パッチ埋め込みを、変形可能畳み込みに基づく変形可能パッチ埋め込みに置換する。
• 変形可能マルチヘッドアテンションと近傍マルチヘッドアテンションを交互に用いる、時空間的に動的な自己注意ブロックを使用する。
• 自己注意の不規則な格子をモデル化するための多尺度変形可能位置エンコーディングを導入する。
• スキップ接続と深い監視を備えた U-Net 風のエンコーダ−デコーダを採用する。
• 深い監視の有無による AgileFormer-T および AgileFormer-B のバリアントを比較する。
• Synapse（2D マルチオーガン）、ACDC（2D 心臓）、Decathlon 脳腫瘍（3D）データセットで評価する。

実験結果

リサーチクエスチョン

• RQ1 時空間的に動的な要素は、医用画像の異なる臓器の形状とサイズの不均一性を ViT-UNets がより適切に扱えるようにするか？
• RQ2 変形可能パッチ埋め込みと多尺度変形可能位置エンコーディングは、固定ウィンドウの ViT-UNets と比べて分割精度とスケーラビリティを向上させるか？
• RQ3 複数のデータセットにおいて、他の最先端 ViT-UNet と比較した場合、モデルサイズとともに AgileFormer の性能はどのようにスケールするか？
• RQ4 AgileFormer のセグメンテーション性能に対する深い監視の影響はどの程度か？
• RQ5 提案された各コンポーネントは、2D および 3D の医用セグメンテーションタスクの両方でどの程度性能を発揮するか？

主な発見

• AgileFormer は 3 データセット（Synapse、ACDC、Decathlon）すべてで最近の SOTA 手法を上回った。
• AgileFormer-T から AgileFormer-B へのモデルスケーリングは、パラメータ数と FLOPs の増加を抑えつつ顕著な DSC 増加をもたらした。
• アブレーション研究は変形可能パッチ埋め込みと MS-DePE が性能向上に寄与することを示した。
• 時空間的に動的なアテンション（DMSA/NMSA）と変形可能コンポーネントは、異なる形状・サイズの対象の取り扱いを改善する。
• 深い監視はさらに DSC の著しい改善をもたらす。
• 3D 実験では AgileFormer-T+DS が 3D ViT-UNets に対して競争力のある結果を示した。