Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining

要約

正確な医用画像セグメンテーションには、局所的な特徴から大域的な依存関係に至るマルチスケール情報の統合が要求される。しかし、畳み込みニューラルネットワーク（CNN）はその局所的な受容野に制約され、視覚変換器（ViT）はその注意メカニズムの高い二次的複雑性に悩まされるなど、既存の手法では長距離の大域的情報をモデル化することは困難である。近年、マンバベースのモデルは、長いシーケンスのモデリングにおいて、その素晴らしい能力から大きな注目を集めている。いくつかの研究により、これらのモデルは様々なタスクにおいて一般的な視覚モデルを凌駕し、より高い精度、より少ないメモリ消費量、より少ない計算負荷を提供できることが実証されている。しかし、既存のMambaベースのモデルは、ほとんどがゼロから学習されたものであり、データ効率の良い医用画像解析に非常に有効であることが証明されている事前学習の力を探求していない。本論文では、ImageNetベースの事前学習の利点を活用し、医療画像セグメンテーションタスクのために特別に設計された、新しいMambaベースのモデル、Swin-UMambaを紹介する。我々の実験結果は、Mambaベースのモデルの性能を向上させる上で、ImageNetベースの学習が重要な役割を果たすことを明らかにしている。Swin-UMambaは、CNN、ViT、および最新のMambaベースのモデルと比較して、大きなマージンをもって優れた性能を示す。特に、AbdomenMRI、Encoscopy、Microscopyの各データセットにおいて、Swin-UMambaは最も近いU-Mambaを平均3.58%のスコアで上回った。Swin-UMambaのコードとモデルはhttps://github.com/JiarunLiu/Swin-UMamba で公開されている。

要約(オリジナル)

Accurate medical image segmentation demands the integration of multi-scale information, spanning from local features to global dependencies. However, it is challenging for existing methods to model long-range global information, where convolutional neural networks (CNNs) are constrained by their local receptive fields, and vision transformers (ViTs) suffer from high quadratic complexity of their attention mechanism. Recently, Mamba-based models have gained great attention for their impressive ability in long sequence modeling. Several studies have demonstrated that these models can outperform popular vision models in various tasks, offering higher accuracy, lower memory consumption, and less computational burden. However, existing Mamba-based models are mostly trained from scratch and do not explore the power of pretraining, which has been proven to be quite effective for data-efficient medical image analysis. This paper introduces a novel Mamba-based model, Swin-UMamba, designed specifically for medical image segmentation tasks, leveraging the advantages of ImageNet-based pretraining. Our experimental results reveal the vital role of ImageNet-based training in enhancing the performance of Mamba-based models. Swin-UMamba demonstrates superior performance with a large margin compared to CNNs, ViTs, and latest Mamba-based models. Notably, on AbdomenMRI, Encoscopy, and Microscopy datasets, Swin-UMamba outperforms its closest counterpart U-Mamba by an average score of 3.58%. The code and models of Swin-UMamba are publicly available at: https://github.com/JiarunLiu/Swin-UMamba

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL