USEMA: a Scalable Efficient Mamba Like Attention for Medical Image Segmentation

技術概要：USEMA – 医用画像セグメンテーションのためのスケーラブルで効率的な Mamba 型アテンション

問題定義
正確な医用画像セグメンテーションは診断パイプラインにとって不可欠ですが、限られたトレーニングデータや複雑で重なり合う物体の幾何学形状といった独自の課題に直面しています。ビジョントランスフォーマー（ViT）は自己アテンションを通じてグローバルな依存関係を捉える能力を提供しますが、その医用画像への応用は以下の 2 つの主要な問題によって阻害されています：

1. 二次的な計算複雑性：純粋な自己アテンションはシーケンス長に対して $O(n^2)$ のコストを課し、高解像度の医用画像では計算上実行不可能となります。
2. アテンションの分散：シーケンス長が増加するにつれて、アテンション行列（softmax 出力）のスコアはゼロに向かって均一に分散する傾向があります。この現象は弱大数の法則に基づいて理論的に裏付けられており、モデルが重要なキーと値を選択する能力を阻害し、アテンション機構を実質的に均一な平均化プロセスに低下させます。

手法
著者らは、Scalable Efficient Mamba-like Attention（SEMA）と呼ばれる新しいアテンション機構を統合したハイブリッド UNet アーキテクチャであるUSEMAを提案します。この手法は、以下の構成要素を通じて、純粋なトランスフォーマーと既存の Mamba ベースモデルの両方の限界に対処します：

• SEMA 機構：SEMA は、線形複雑性を維持しつつ完全な自己アテンションの挙動を模倣するために、局所近似とグローバル近似を組み合わせます。
  • 局所コンポーネント：ウィンドウアテンションを利用してトークンの局所化と焦点を維持し、長いシーケンスに固有の分散問題を回避します。
  • グローバルコンポーネント：二次的なコストなしにグローバルな文脈を捉えるため、SEMA は理論的に整合性のある算術平均を採用します。この項は分散不等式で観測される $O(1/n)$ の境界を反映し、アテンションのグローバルな側面を近似します。グローバル項は、すべての値ベクトルの平均として定義され、シーケンス全体にブロードキャストされます。
  • 統合：SEMA ブロックは、Mamba の条件付き位置エンベディング、レイヤー正規化、ゲート機構（SiLU 活性化や深度方向畳み込みなど）に類似したメカニズムを利用する Mamba マクロ構造に挿入されます。
• USEMA アーキテクチャ：このモデルは対称的な U 字型構造（UNet）を採用します。
  • エンコーダ：残差ブロックに続いて SEMA ブロックで構成されます。特徴量はアテンション機構に入る前に $(B, C, H, W)$ から $(B, C, HW)$ に再整形されます。回転位置エンベディングと局所的に強化された位置符号化が、空間認識を強化するために使用されます。
  • ボトルネック：空間圧縮によりシーケンス長が十分に小さく計算的に管理可能であるため、ボトルネックではウィンドウアテンションではなく完全な自己アテンションが利用されます。
  • デコーダ：残差ブロックと転置畳み込みで構成され、エンコーダからの階層的特徴を融合させるためにスキップ接続を利用します。

主要な貢献

1. ハイブリッドアテンション設計：本論文は、分散を防ぐために局所ウィンドウアテンションを補完するものとして、算術平均を自己アテンションのグローバル近似として使用する理論的根拠を示す SEMA を導入します。
2. アーキテクチャ統合：USEMA は、UNet フレームワーク内で CNN の局所特徴抽出効率と Mamba 型アテンションのグローバルモデリング能力を成功裡に融合させます。
3. 包括的なベンチマーク：このモデルは、解像度やセグメンテーション対象（臓器、器具、細胞のセグメンテーション）が異なる 3 つの多様な医用画像モダリティ（腹部 MRI、内視鏡、顕微鏡）で評価されました。

実験結果
USEMA は、3 つのデータセットにおいて、最先端のトランスフォーマーベースモデル（UNETR, Swin-UNETR, nnFormer）および Mamba ベース/インスパイアモデル（U-Mamba, Swin-UMamba, Mamba UNet, MLLA-UNet）とベンチマーク比較されました：

• 腹部 MRI（MICCAI 2022 AMOS）：USEMA は、Dice 類似係数（DSC）で0.7704、正規化表面距離（NSD）で0.8345を達成しました。これは、最良のトランスフォーマーベースライン（nnFormer）を DSC で 5.83%、最良の Mamba ベースライン（U-Mamba Enc）を DSC で 10.53% 上回ります。注目すべきは、USEMA が52M パラメータでこれらの結果を達成したことで、これは最高性能のトランスフォーマーモデルより 13% 少ないパラメータ数です。
• 内視鏡（MICCAI 2017）：USEMA は DSC で0.6463、NSD で0.6621を達成し、トランスフォーマーベースラインを DSC で約 5.6%、Mamba ベースラインを DSC で約 13.5% 上回りました。
• 顕微鏡（NeurIPS 2022）：このインスタンスセグメンテーションタスクにおいて、USEMA は F1 スコアで0.5791を達成し、最良のトランスフォーマーベースラインを 8.61%、Mamba インスパイアの MLLA-UNet を 19.23% 上回りました。また、U-Mamba Enc よりも 43% 少ないパラメータ数を利用しました。

アブレーション研究
アブレーション研究により、グローバルアテンション近似の必要性が確認されました。算術平均項を除去し（ウィンドウアテンションのみを使用）、腹部 MRI データセットの DSC で 0.013 の低下など、すべてのデータセットで性能低下が見られ、グローバル項の寄与が実証されました。

意義と主張
著者らは、USEMA が計算効率とセグメンテーション精度のバランスを効果的に取っていると主張しています。純粋な自己アテンションの分散特性と Mamba の線形複雑性を活用することで、USEMA は以下の点でスケーラブルな解決策を提供します：

• 多様な 2 次元医用画像ベンチマークにおいて、純粋な畳み込みモデル、トランスフォーマーベースモデル、Mamba ベースモデルを上回る性能を発揮します。
• 完全な自己アテンショントランスフォーマーと比較して計算コストを削減します。
• 理論的に裏付けられた効率的なグローバル近似（算術平均）が、医用画像分析において局所アテンション機構を成功裡に補完し得ることを実証します。

本論文は、平均化メカニズムを学習済みまたは重み付き平均化へ拡張し、USEMA を大規模な病理画像へ適用する計画で締めくくられていますが、これらの特定の将来の分野における現在の成功については主張していません。