技術的サマリー：Patch-MoE Mamba

問題提起

医療画像セグメンテーションは、畳み込みニューラルネットワーク（CNN）とトランスフォーマーに大きく依存しています。しかし、CNN は局所的な受容野に本質的に制限されており、長距離依存性のモデル化に苦慮する一方、トランスフォーマーはしばしば二次的な計算量とメモリ量というコストを伴います。最近、状態空間モデル（SSM）、特に Mamba ベースのアーキテクチャが、線形シーケンシャル複雑性を持つ効率的な代替手段として登場しています。

それらの可能性にもかかわらず、既存の Mamba ベースのセグメンテーションモデルは、2 つの重要な限界に直面しています：

1. 局所空間構造の破壊： ほとんどのモデルは、ピクセル単位のラスタースキャンを用いて 2 次元特徴マップを 1 次元シーケンスに変換します。このプロセスは、シーケンス内で空間的に隣接するピクセルを分離し、局所的な 2 次元空間の整合性を弱体化させる可能性があります。これは、正確な境界局所化と局所構造の保持を必要とする密な予測タスクにとって有害です。
2. 柔軟性に欠ける方向性融合： 現在の手法は、通常、単純な加算によって複数の走査方向からの出力を融合します。この固定された集約は、すべての空間位置においてすべての方向が等しく重要であると仮定しており、医療ターゲットの多様なサイズ、形状、テクスチャ、および境界の複雑さ（例えば、小さくコントラストの低い病変対、大きく不規則な構造）に適応できません。

手法

著者らは、前述の限界に対処するために設計された、パッチ順序の混合エキスパート（MoE）状態空間アーキテクチャであるPatch-MoE Mambaを提案します。このアーキテクチャは U-Net 風のデザインに従い、VM-UNetV2 における標準的なビジュアル状態空間（VSS）ブロックを、新しいPatch-MoE VSS ブロックに置き換えています。

1. 階層的パッチ順序走査

個々のピクセルをラスタースキャンする代わりに、この手法は特徴マップを局所パッチに分割し、パッチ順序で方向走査を行います。

• メカニズム： 特徴マップが与えられると、空間グリッドは $p \times p$ サイズの非重なりパッチに分割されます。各パッチ内では、次のパッチへ移動する前に、ピクセルが行優先順序で列挙されます。
• 利点： これにより、パッチ内のすべてのピクセルがシーケンス内で連続して保持され、パッチ間での長距離依存性のモデル化を可能にしながら、局所的な空間近傍を保持します。
• 階層的戦略： 多スケールの文脈を捉えるために、モデルはパッチサイズの集合 $\{p^{(1)}, p^{(2)}, \dots\}$ を採用します。より小さなパッチは微細な境界の詳細を捉え、より大きなパッチは粗い構造をモデル化します。異なる走査方向（前方、後方、幅 - 高さ前方、幅 - 高さ後方）は、異方性構造を処理するために異なるパッチサイズを利用できます。

2. MoE ベースの方向性融合

固定された加算に代わるために、著者らは方向性出力を「エキスパート」として扱う適応的融合モジュールを導入します。

• エキスパートの構築： このモジュールは 5 つのエキスパートを利用します：
  • 4 つの方向固有エキスパート（$E_1$ から $E_4$）：4 つの方向性 Mamba スキャナーの正規化された出力。
  • 1 つの結合エキスパート（$E_{concat}$）：4 つの正規化マップをチャネル次元に沿って結合し、$1\times1$ 畳み込みを通じて投影して形成されます。このエキスパートは方向間相関を捉え、方向非依存の包括的な視点を提供します。
• 空間認識ゲート制御： ルーターネットワークが位置依存の重みを計算し、これらのエキスパートを適応的に融合します。ルーターは、局所的文脈（深度方向畳み込みを介して）とグローバルな文脈（適応平均プーリングを介して）のブレンドによって駆動され、学習可能なスカラー $\alpha$ によって制御されます。
• 残差安定化： 残差バイパスが、融合結果に生の方向性出力を追加します（$Z_l = \hat{Y}_l + \sum Y_l^{(i)}$）。これにより、堅牢なベースライン応答が保証され、強力な方向性信号が保持され、トレーニングが安定します。

3. 意味情報と詳細の注入（SDI）

エンコーダーの出力は、SDI モジュール（U-Net v2 から適応されたもの）によって処理され、上位レベルからの意味情報を注入し、ハダマール積を介して下位レベルからのより詳細な情報を統合することで、特徴マップを強化します。

主要な貢献

1. パッチ順序走査： ピクセル単位のラスタースキャンをパッチ順序のトラバーサルに置き換える新しい走査メカニズム。これにより、状態空間モデルの効率性を維持しながら、局所的な 2 次元空間の整合性を効果的に保持します。
2. MoE ベースの方向性融合： 単純な加算を学習可能なゲート制御ネットワークに置き換える適応的融合モジュール。これは、空間的文脈に基づいて複数の方向エキスパートと結合エキスパートを動的に重み付けし、モデルが多様な病変スケールと複雑さに対応できるようにします。
3. 階層的多スケールモデリング： 階層的パッチサイズの統合により、ネットワークは微細な境界の詳細と粗い意味構造を同時に捉えることができます。
4. 残差安定化： 生の方向性出力に対する残差経路の組み込みにより、トレーニングの安定性が保証され、重要な方向性信号の損失が防止されます。

実験結果

著者らは、Patch-MoE Mamba を 5 つの公開ポリープセグメンテーションデータセット（Kvasir-SEG、ClinicDB、ColonDB、ETIS、CVC-300）と 2 つの皮膚病変セグメンテーションデータセット（ISIC 2017、ISIC 2018）で評価しました。

• 性能： Patch-MoE Mamba はすべてのデータセットで最先端の性能を達成し、CNN ベース（U-Net、U-Net v2）および Mamba ベースのベースライン（VM-UNet、VM-UNetV2）を一貫して上回りました。
  • 低コントラストと小さなポリープで知られる挑戦的なETISデータセットにおいて、モデルは Dice スコアを 72.56%（VM-UNetV2）から**74.04%**に向上させました。
  • ISIC 2018では、Dice スコア**89.34%**を達成し、以前の最高値 88.36% を上回りました。
• アブレーション研究：
  • ラスタースキャンをパッチ順序走査に置き換えるだけで、平均 Dice は 83.46% から 84.02% に向上しました。
  • MoE 融合モジュールを追加することで、平均はさらに**84.30%**に引き上げられました。
  • 階層的パッチサイズ（初期段階では粗く、深い段階では細かく）が最良の結果をもたらしました。
• 計算複雑性： パッチ順序走査自体は計算コストを増加させません（トークン数は変化しません）が、結合エキスパートの組み込みにより、パラメータ数は約 70M、FLOPs は約 28G に増加します。著者らは、このオーバーヘッドは MoE メカニズムに必要な高密度なチャネル相互作用に起因すると指摘しています。

意義と主張

本論文は、Patch-MoE Mamba が医療画像セグメンテーションのための一般的かつ効果的な解決策を提供すると主張しています。既存の Mamba モデルの特定の限界、すなわち局所空間構造の喪失と方向性融合の硬直性に対処することで、提案されたアーキテクチャは、サイズ、形状、境界の複雑さが異なる多様な医療ターゲットに対する優れた適応性を示しています。結果は、パッチ順序走査による局所空間近傍の保持と、適応的エキスパート融合の組み合わせが、特に低コントラストまたは小さな病変を含む困難なシナリオにおいて、境界局所化とセグメンテーション精度を大幅に向上させることを示唆しています。