[論文レビュー] ViM-UNet: Vision Mamba for Biomedical Segmentation

Anwai Archit, Constantin Pape|arXiv (Cornell University)|Apr 11, 2024

Image Processing Techniques and Applications被引用数 4

ひとこと要約

この論文では、従来のUNetおよびUNETRネットワークに代わる、ビジョン・マンバーアーキテクチャに基づく新しい生体医療画像セグメンテーションモデル、ViM-UNetを紹介する。LIVECellおよびCREMIの2つの顕微鏡インスタンスセグメンテーションタスクにおいて、競争力あるか、それ以上の性能を達成している。UNetRに比べて顕著にパラメータ数が少なく、高速であるため、文脈を多く必要とする生体医療タスクにおけるマンバーベースモデルの可能性を示している。

ABSTRACT

CNNs, most notably the UNet, are the default architecture for biomedical segmentation. Transformer-based approaches, such as UNETR, have been proposed to replace them, benefiting from a global field of view, but suffering from larger runtimes and higher parameter counts. The recent Vision Mamba architecture offers a compelling alternative to transformers, also providing a global field of view, but at higher efficiency. Here, we introduce ViM-UNet, a novel segmentation architecture based on it and compare it to UNet and UNETR for two challenging microscopy instance segmentation tasks. We find that it performs similarly or better than UNet, depending on the task, and outperforms UNETR while being more efficient. Our code is open source and documented at https://github.com/constantinpape/torch-em/blob/main/vimunet.md.

研究の動機と目的

• CNNベースのUNetおよびトランスフォーマー基盤のUNETRに起因する制限、特に高いパラメータ数と計算コストを解消すること。
• 最近提案されたグローバルな受容 field と効率的なシーケンスモデリングを有するビジョン・マンバー（ViM）が、生体医療画像セグメンテーションにおいてトランスフォーマーの優れた代替手段となり得るかを検討すること。
• 神経線維のセグメンテーションなど、大規模な文脈的理解を要する電子顕微鏡画像における挑戦的な顕微鏡データセットにおいて、ViM-UNetの性能を評価すること。
• 異なるデータセットにおいて、UNet、UNETR、nnUNet、U-Mambaと比較して、ViM-UNetの正確性、推論速度、メモリ効率を評価すること。
• パラメータ数と推論時間の分析を通じて、特にデータ量が少ない状況や3Dセグメンテーションのシナリオにおいて、モデルの汎化性能と効率性を検証すること。

提案手法

• Vision Mamba（ViM）エンコーダーを、双方向の状態空間モデル（SSM）を用いて適応させ、UNETRのViTバックボーンを置き換えることで、計算コストを低減しつつグローバルな文脈モデリングを可能にする。
• ビジョン・トランスフォーマーと同様に16×16ピクセルのパッチ処理を採用するが、多頭部アテンションの代わりにSSMを用いることで、長距離依存関係を効率的に捉える。
• すべてのモデルで公平な比較が行えるよう、標準的なUNetスタイルのデコーダーを採用。各ヘッドに2つの畳み込み層と逆畳み込みを用い、同一のアーキテクチャを共有する。
• すべてのモデルを、初期学習率1e-4で最適化し、100kイテレーションにわたり学習率の低下をプレートオーブンで実行するAdam最適化法で訓練する。
• 2つのインスタンスセグメンテーションヘッドを実装：1つは前景および境界の確率マップ（ウォーターシェド後処理を伴う）、もう1つは中心点および距離予測。
• LIVECell（位相差顕微鏡）およびCREMI（ボリュームEM、2Dスライス）でモデルを評価し、主に平均セグメンテーション正確性を指標とする。

Figure 1: Example images with segmentation from ViM-UNet Small .

実験結果

リサーチクエスチョン

• RQ1 ビジョン・マンバー基盤のアーキテクチャ（ViM-UNet）は、挑戦的な生体医療インスタンスセグメンテーションタスクにおいて、UNet や UNETR と同等またはそれ以上のセグメンテーション性能を達成できるか？
• RQ2 ViM-UNetのグローバルな視野が、神経線維のセグメンテーションなど、大規模な文脈理解を要するタスクにおいて、UNet よりも性能優位性を提供するか？
• RQ3 特にUNETRの高いパラメータ数と計算要求を踏まえると、ViM-UNetは推論速度およびメモリ効率においてUNETRに比べて優れているか？
• RQ4 豊富な事前学習を必要とせず、ViM-UNetは多様なデータセット、特に異なる形態的特徴を有するデータセットにおいても汎化性能に優れているか？
• RQ5 文脈が極めて重要な3Dセグメンテーションや細胞追跡の文脈において、ViM-UNetはトランスフォーマー基盤のモデルの実用的で効率的な代替手段となり得るか？

主な発見

• LIVECellデータセットでは、ViM-UNetは距離予測付きのUNetと同等の性能（0.05秒の推論時間）を達成し、パラメータ数が多いにもかかわらずUNETR（0.54秒）を上回った。
• LIVECellにおける境界セグメンテーションでは、ViM-UNetはUNetに次いで2位となり、UNETRが最悪の性能を示したのとは対照的だった。
• CREMIデータセットでは、ViM-UNetが最高の平均セグメンテーション正確性を記録し、UNetおよびUNETRを上回った。これは、大規模構造のタスクにおいてグローバルな文脈モデリングが大きな利点をもたらすことを示している。
• ViM-UNetは、Tinyで18Mパラメータ、Smallで39Mパラメータにとどまり、UNETRの113M（Base）、334M（Large）、665M（Huge）に比べ顕著に少ない。しかし、競争力のある性能を維持している。
• 推論時間はViM-UNetがTinyとSmallでそれぞれ0.05秒と非常に短く、UNETRの0.15秒から0.54秒に比べ顕著に短縮された。また、訓練には最大48GBのVRAMを要するUNETRとは異なり、ViM-UNetは10GB未満のVRAMで十分だった。
• これらの結果から、ViM-UNetは、アテンション機構のオーバーヘッドがなく、効率的なグローバルモデリングを実現しているため、3Dセグメンテーションや細胞追跡など、文脈を多く必要とするタスクにおいて特に効果的であると考えられる。

Figure 2: Results for our and external methods; circles highlight the three best methods.