Hybrid Transformer-Mamba for Weakly Supervised Volumetric Medical Segmentation

技術要約：弱教師あり学習による体積的医療画像セグメンテーションのためのハイブリッド・トランスフォーマー・マンバ（TranSamba）

問題提起
体積的な医療画像セグメンテーション（例：CT、MRI）は、通常、労力を要するボクセルレベルのアノテーションを必要とする。弱教師ありセグメンテーションは、局在情報を持たないスライスレベル（平面レベル）のラベルのみを用いてモデルを学習させることで、この解決策を提示する。Class Activation Mapping (CAM) は分類モデルから粗い局在化を抽出できるが、既存の手法は医療データの体積的な性質により、重大な課題に直面している。現在の手法の多くは、個々のスライスを独立して処理する2Dエンコーダに依存しており、3Dコンテキストを無視している。さらに、標準的なクロスプレーン自己注意（Self-Attention, SA）を通じて体積的コンテキストを捉えようとすると、計算量が平面数（$N$）に対して二次関数的に増大するため、計算コストが極めて高くなる。

手法：TranSamba
著者らは、Vision Transformer (ViT) と Mamba (Structured State Space Model) アーキテクチャを組み合わせたハイブリッド構造である TranSamba を提案し、これらの制限に対処する。コアとなる設計思想は、Mamba の効率性をクロスプレーンの相互作用に活用しつつ、インプレーン（平面内）の物体局在化における Transformer の有効性を維持することである。

• アーキテクチャの概要: エンコーダは $L$ 個のハイブリッド層で構成される。各層には Cross-Plane Mamba (CPM) ブロックと、それに続く Transformer ブロックが含まれる。
• クロスプレーン・モデリング (CPM): 3Dコンテキストを効率的に捉えるため、$N$ 個の隣接する平面からのパッチ・トークンをスタックし、$M \times N$ の長さ（$M$ は平面あたりのパッチ数）の単一のシーケンスへと再構成する。CPM ブロックは Mamba モジュールを使用して、これらのトークン間の相互作用をモデル化する。重要な点として、クラス・トークンは分離され、Mamba ブロックから除外される。これにより、平面のグローバルな表現を保持する。これにより、モデルは線形時間で平面間の依存関係を学習することができる。
• インプレーン・モデリング (Transformer): CPM ブロックの出力は、再び $N$ 個の並列なシーケンスへと再構成される。これらは、Class-to-Patch (C2P) 自己注意を用いた標準的な Transformer ブロックによって処理される。クロスプレーン・ステップによって強化されたトークン表現は、インプレーンの注意マップを向上させ、物体局在化を改善する。
• 学習と推論: モデルは平面レベルのラベルを用いて学習される。損失関数は、クラス・トークンのグローバル平均プーリング (GAP) と、パッチ・トークンのグローバル加重ランキングプーリング (GWRP) から導出されるクロスエントロピー項を組み合わせたものである。推論時には、すべての Transformer ブロックからの C2P 注意重みを集計することで、注意マップが生成される。

主な貢献

1. ハイブリッド・アーキテクチャ: CPM ブロックを ViT バックボーンに統合することで、インプレーンの自己注意の強みを損なうことなく、効率的なクロスプレーン・モデリングを可能にする TranSamba を導入した。
2. 計算効率: 本アーキテクチャは、平面数（$N$）に対して線形時間計算量を達成し、標準的なクロスプレーン自己注意の二次関数的なスケーリングとは対照的に、バッチ処理に対して一定の空間計算量を維持する。
3. 最先端の性能 (SOTA): 本手法は、クロスプレーン・モデリングが、従来の作業で頻繁に使用される複雑なマルチスケール集約や自己教師あり学習の正則化戦略に依存することなく、体積的医療画像セグメンテーションにおける汎用的な戦略であることを示し、多様なモダリティと病理において優れた結果を達成した。

実験結果
TranSamba は、以下の 3 つの異なるデータセットで評価された：

• BraTS: 脳腫瘍セグメンテーション (T2 FLAIR)。

• KiTS: 腎臓腫瘍セグメンテーション。

• LASC: 左心房腔セグメンテーション。

• 性能: TranSamba はこれら 3 つのデータセットすべてで SOTA の結果を達成し、2 番目に優れた手法を、Dice 類似係数 (DSC) で平均 5.3%、Intersection-over-Union (IoU) で 2.4% 上回った。

• アブレーション研究:

  • クロスプレーン vs インプレーン: CPM ブロックをインプレーン Mamba ブロック (V2) に置き換えると性能が大幅に低下した。これは、改善が単なる Mamba の使用によるものではなく、具体的にクロスプレーン・モデリングに由来することを裏付けている。
  • モデリング範囲: 16 平面間でのモデリングが最良の結果を示した。これは、連続的な構造に対しては、長距離モデリングよりも局所的なクロスプレーン相互作用の方が有益であることを示唆している。
  • レイヤー設計: CPM ブロックを Transformer ブロックの「前」に配置する設計 (CrossIn デザイン) が最適であった。これにより、注意の抽出が強化されたクロスプレーン・コンテキストの恩恵を受けることが保証される。

• 物体サイズ分析: 性能は物体サイズと正の相関を示した。TranSamba は大型および巨大な物体に対しては優れていたが、小さな物体（医療セグメンテーションにおける既知の問題）については課題があり、特に物体が小さい傾向にある KiTS データセットにおいてそれが顕著であった。

意義と主張
本論文は、TranSamba が 3D コンテキストの必要性と、計算上の制約の間のギャップを効果的に埋めるものであると主張している。クロスプレーン・モデリングとインプレーン・モデリングを分離することにより、著者らは、効率的な線形時間のクロスプレーン・モデリングが、体積的な医療画像セグメンテーションにおける物体局在化を向上させるための、重要かつ汎用的なコンポーネントであることを実証した。本研究は、将来の臨床応用には後処理やマルチステップのアプローチが必要になる可能性があることを示唆しているが、核となるアーキテクチャの革新は、多様な医療画像シナリオにおける弱教師ありタスクのための堅牢な基盤を提供するものである。