[論文レビュー] VMambaMorph: a Multi-Modality Deformable Image Registration Framework based on Visual State Space Model with Cross-Scan Module

Ziyang Wang, Jian-Qing Zheng+2 more|arXiv (Cornell University)|Apr 7, 2024

Advanced Image and Video Retrieval Techniques被引用数 4

ひとこと要約

VMambaMorphは、3D最適化されたクロススキャンモジュールと再帰的ファインチューニング戦略を備えたハイブリッドVisual Mamba-CNNアーキテクチャを活用する、新規の3Dマルチモodal変形画像登録フレームワークを提案する。公に提供されたMR-CTデータセットにおいて、VoxelMorph、TransMorph、MambaMorphを上回る最先端の登録精度を達成し、Diceスコア82.94±2.01を記録するとともに、計算効率が向上した。

ABSTRACT

Image registration, a critical process in medical imaging, involves aligning different sets of medical imaging data into a single unified coordinate system. Deep learning networks, such as the Convolutional Neural Network (CNN)-based VoxelMorph, Vision Transformer (ViT)-based TransMorph, and State Space Model (SSM)-based MambaMorph, have demonstrated effective performance in this domain. The recent Visual State Space Model (VMamba), which incorporates a cross-scan module with SSM, has exhibited promising improvements in modeling global-range dependencies with efficient computational cost in computer vision tasks. This paper hereby introduces an exploration of VMamba with image registration, named VMambaMorph. This novel hybrid VMamba-CNN network is designed specifically for 3D image registration. Utilizing a U-shaped network architecture, VMambaMorph computes the deformation field based on target and source volumes. The VMamba-based block with 2D cross-scan module is redesigned for 3D volumetric feature processing. To overcome the complex motion and structure on multi-modality images, we further propose a fine-tune recursive registration framework. We validate VMambaMorph using a public benchmark brain MR-CT registration dataset, comparing its performance against current state-of-the-art methods. The results indicate that VMambaMorph achieves competitive registration quality. The code for VMambaMorph with all baseline methods is available on GitHub.

研究の動機と目的

• 従来のディープラーニングベースの画像登録手法が、長距離依存関係を捉えることや、複雑なマルチモダリティ変形を処理することに課題を抱えているのを是正すること。
• 3Dボリューム画像登録における視覚的状態空間モデル（VMamba）とクロススキャンモジュールの応用を検討し、グローバル特徴モデリングを向上させること。
• 3Dマルチモダリティ登録に特化したハイブリッドVMamba-CNN U-Netアーキテクチャを設計し、特徴抽出と変形場学習を強化すること。
• 登録精度を向上させつつ微分同相性の性質を保持する再帰的ファインチューニングおよび推論フレームワークを開発すること。
• 公に提供されたMR-CT登録ベンチマーク上で提案手法を検証し、VoxelMorph、TransMorph、MambaMorphを含む最先端手法と比較すること。

提案手法

• U字型のエンコーダデコーダアーキテクチャを採用し、3D VMambaブロックと再設計された2Dクロススキャンモジュールを統合。3Dボリューム特徴処理に適応させた。
• 空間特徴表現を強化するための細分化された特徴抽出器を導入し、異なるモダリティ間で複雑な解剖学的構造を捉える能力を向上させた。
• ネットワークが予測された変位場をフィードバックとして用いて、反復的に変形場を精緻化する再帰的トレーニングおよび推論戦略を実装した。
• 損失関数には正規化相互相関の不一致度測定と滑らかさ正則化項を組み合わせ、正確で微分同相性の保証された変換を実現した。
• AdamWを用いた最適化と学習率スケジューリングを実施し、公に提供された脳MR-CT登録データセットを用いてエンドツーエンドで学習した。
• 変形場は連続的な変位ベクトル場として予測され、ヤコビアン行列式正則化により微分同相性の性質が保証された。

Figure 1: The Recursive Registration Framework of VMambaMorph.

実験結果

リサーチクエスチョン

• RQ1 クロススキャンモジュールを備えたVisual Mambaアーキテクチャは、3Dマルチモダリティ医療画像登録において、長距離依存関係を効果的にモデリングできるか？
• RQ2 3D最適化VMambaブロックと3D CNNの統合は、標準的なCNNやTransformerベースの手法と比較して、登録精度をどのように向上させるか？
• RQ3 提案された再帰的登録フレームワークは、変形場の品質をどの程度向上させ、登録誤差を低減するか？
• RQ4 細分化された特徴抽出器の導入により、MRとCT画像間の構造的・強度的変動を処理する能力が顕著に向上するか？
• RQ5 標準ベンチマーク上で、VoxelMorph、TransMorph、MambaMorphといった最先端手法と比較して、VMambaMorphは性能と効率性の両面で優れているか？

主な発見

• VMambaMorphは、公に提供されたMR-CT登録ベンチマークでDiceスコア82.94±2.01を達成し、VoxelMorph（62.42±3.29）、TransMorph（78.46±2.13）、MambaMorph（82.49±1.99）を上回った。
• 95%ヒューリスト距離が1.35±0.18 mmまで低下し、ベースラインモデルと比較して優れた空間整合性の精度を示した。
• 非正のヤコビアン行列式の割合が1.04±0.05%に抑えられ、変形場における微分同相性の性質が保持されていることが確認された。
• K=2の再帰的推論戦略が、精度と推論時間のバランスを最良に保ち、実行時間0.19秒、GPUメモリ使用量3.93 GBを記録した。
• アブレーションスタディにより、細分化された特徴抽出器の導入により、ベースラインと比較してDiceスコアが最大4.5%向上した。
• VMambaMorphは、VoxelMorph、TransMorph、MambaMorphと比較して、トレーニング損失がより早く安定し、より低い値に収束した。これは、収束性と一般化性能に優れていることを示している。

Figure 2: The Architecture of VMambaMorph, and the Details of Visual State Space Block.