[论文解读] ViM-UNet: Vision Mamba for Biomedical Segmentation

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

Image Processing Techniques and Applications被引用 4

一句话总结

该论文提出 ViM-UNet，一种基于视觉Mamba架构的新型生物医学图像分割模型，取代了传统的UNet和UNETR网络。在两项显微镜实例分割任务（LIVECell 和 CREMI）中，该模型实现了具有竞争力或更优的性能，同时相比 UNETR 显著减少了参数量并提升了推理速度，展示了基于Mamba的模型在上下文密集型生物医学任务中的潜力。

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和基于Transformer的UNETR在生物医学分割中的局限性，特别是高参数量和高计算成本。
• 探究近期提出的具有全局感受野和高效序列建模能力的视觉Mamba（ViM）架构，是否可作为生物医学图像分割中Transformer的更优替代方案。
• 在需要大范围上下文理解的挑战性显微镜数据集上评估ViM-UNet的性能，例如电子显微镜中的神经元突起分割。
• 在多种数据集上，对比ViM-UNet与UNet、UNETR、nnUNet和U-Mamba在精度、推理速度和内存效率方面的表现。
• 通过分析参数量和推理时间，验证模型在低数据量场景和3D分割任务中的泛化能力与效率。

提出的方法

• 将视觉Mamba（ViM）编码器与双向状态空间模型（SSMs）结合，替代UNETR中的ViT主干网络，实现在更低计算成本下对全局上下文的建模。
• 采用类似视觉Transformer的分块处理方式（16×16分块），但用SSMs替代多头注意力机制，以高效实现长距离依赖。
• 采用标准的UNet风格解码器，每个头包含两层卷积和转置卷积，所有模型保持相同架构以确保公平比较。
• 所有模型均使用Adam优化器训练，初始基础速率设为1e-4，并在验证损失不再下降时采用学习率衰减策略，共训练100k轮。
• 实现两个实例分割头：一个用于前景和边界概率图（经分水岭后处理），另一个用于中心点和距离预测。
• 在LIVECell（相差显微镜）和CREMI（体积电子显微镜，2D切片）数据集上评估模型，主要指标为平均分割准确率。

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

实验结果

研究问题

• RQ1 基于视觉Mamba的架构（ViM-UNet）是否能在具有挑战性的生物医学实例分割任务中，实现与UNet和UNETR相当或更优的分割性能？
• RQ2 ViM-UNet的全局感受野是否在需要大范围上下文理解的任务（如电子显微镜图像中的神经元突起分割）中，相比UNet提供性能优势？
• RQ3 在推理速度和内存效率方面，ViM-UNet与UNETR相比如何，尤其考虑到UNETR存在高参数量和高计算需求？
• RQ4 ViM-UNet是否能在无需大量预训练的情况下超越UNETR，并在具有不同形态特征的多样化数据集中展现出良好的泛化能力？
• RQ5 在上下文至关重要的3D分割和细胞追踪任务中，ViM-UNet是否是一种可行且高效的Transformer基模型替代方案？

主要发现

• 在LIVECell数据集上，ViM-UNet在使用距离预测时达到与UNet相当的性能（推理时间0.05秒），并优于UNet（0.54秒），尽管UNet参数量更大。
• 在LIVECell的边界分割任务中，ViM-UNet排名第二，仅次于UNet，且显著优于表现最差的UNet。
• 在CREMI数据集上，ViM-UNet实现了最高的平均分割准确率，优于UNet和UNETR，表明其全局上下文建模在大结构任务中具有显著优势。
• ViM-UNet仅使用18M参数（Tiny）或39M（Small），远低于UNet的113M（Base）、334M（Large）和665M（Huge），同时保持了具有竞争力的性能。
• ViM-UNet的推理时间显著更低（Tiny和Small均为0.05秒），而UNet的推理时间在0.15秒至0.54秒之间；且训练仅需≤10GB显存，而UNet最高需48GB显存。
• 结果表明，由于其高效的全局建模能力且无需注意力机制的开销，ViM-UNet在需要大范围上下文的任务（如3D分割和细胞追踪）中尤为有效。

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