MambaMIM: 预训练Mamba与状态空间标记插值

MambaMIM: Pre-training Mamba with State Space Token-interpolation

GitHub - FengheTan9/MambaMIM   |   2408.08070 (arxiv.org)

Abstract

Generative self-supervised learning demonstrates outstanding representation learning capabilities in both Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). However, there are currently no generative pre-training methods related to selective state space models (Mamba) that can handle long-range dependencies effectively. To address this challenge, we introduce a generative self-supervised learning method for Mamba (MambaMIM) based on Selective Structure State Space Sequence Token-interpolation (S6T), a general-purpose pre-training method for arbitrary Mamba architectures. Our method, MambaMIM, incorporates a bottom-up 3D hybrid masking strategy in the encoder to maintain masking consistency across different architectures. Additionally, S6T is employed to learn causal relationships between the masked sequence in the state space. MambaMIM can be used on any single or hybrid Mamba architectures to enhance the Mamba long-range representation capability. Extensive downstream experiments reveal the feasibility and advancement of using Mamba for pretraining medical image tasks.

Introduction

近年来，状态空间模型（尤其是其变体Mamba）在长序列建模中展现了高效性和有效性，通过选择性扫描和硬件感知实现，Mamba在自然语言和DNA序列建模方面与Transformer竞争激烈。在计算机视觉领域，Vision Mamba作为一种潜在的解决方案，有效处理图像的长距离依赖，并在广泛的视觉任务中表现出色。当从更广泛的计算机视觉应用转向医学图像分析（MIA）时，Vision Mamba也表现出竞争力。然而，由于MIA中标签数据的稀缺性以及状态空间序列的固有属性，训练高性能的纯Vision Mamba以完成医学下游任务具有挑战性。混合架构可以利用卷积的归纳偏置来弥补局部特征提取能力的不足，但仍难以突破性能瓶颈。

幸运的是，自监督学习为MIA提供了一种有前景的方法，通过在大规模未标注图像集上预训练模型，然后迁移到下游任务，这种预训练-微调范式可以显著提升数据稀缺的下游任务性能。在视觉任务的自监督学习方法中，对比学习方法和生成学习方法最为成功。其中，以一系列掩码图像建模（MIM）方法为代表的生成学习方法展现出更强的迁移能力。然而，MIM方法需要处理掩码和可见块，这通常涉及为特定架构定制的策略，使得现有MIM方法难以直接应用于Vision Mamba模型。

如何在生成式风格下预训练Vision Mamba模型仍是一个具有科学价值但尚未实现的研究课题，主要原因是将3D空间体积建模为状态空间序列的难度较大。在预训练Mamba模型时采用掩码图像建模（MIM）的策略，一个直观的想法是丢弃掩码块并在掩码位置插入可学习标记以重建掩码像素，这一技术之前在ViTs的MAE方法中使用过。然而，初步实验表明，尽管可学习标记对纯Mamba和混合Mamba有一定效果，但其性能仍不满意。这一不足归因于可学习标记的插入与因果和输入依赖的选择性扫描属性不符，因为这些标记忽略了状态空间内的结构序列关系，导致历史信息被随机初始化的可学习标记选择，可能无法有效传递给后续标记。虽然CNN-Mamba混合架构可以在一定程度上提升性能，但上述问题依然存在。此外，对于CNN-Mamba混合模型，在卷积层和Mamba层之间保持掩码一致性对于端到端训练至关重要，不同架构间的不一致掩码可能导致像素分布偏移，对表征学习产生负面影响。

为解决上述挑战，本文提出了一种针对Mamba的掩码建模策略MambaMIM，用于在大规模3D CT数据集上预训练CNN-Mamba混合模型。作为预训练模型，MambaMIM采用了一种独特的标记生成方法——选择性结构状态空间序列标记插值（S6T），该方法能有效利用Mamba模型中结构状态空间序列之间的关系进行标记生成。在编码器中，本文通过自下而上的掩码建模来保持掩码一致性，从而确保CNN-Mamba的端到端预训练。值得注意的是，MambaMIM不仅限于预训练混合架构，还可用于预训练纯Vision Mamba模型（混合模型中的特例），如图2所示。通过本文提出的MambaMIM，成功地将基于生成的预训练方法应用于纯Mamba和混合Mamba架构。据本文所知，MambaMIM是首个专为Mamba和选择性结构状态空间序列设计的自监督预训练方法。在多个医学图像分割任务中，使用MambaMIM预训练的CNN-Mamba混合模型表现优于其他最先进的自监督预训练方法和架构。

Figure 2: Vanilla Mamba v.s. Hybrid Mamba with different mask token strategies pre-trained on the BTCV dataset for the 3D segmentation task. vanilla Mamba: Vision Mamba replaces ViT in UNETR; hybrid Mamba (HyMamba): MedNeXt (CNN) + Mamba. The improvements brought by S6T surpass previous SSL-methods with the learnable token, and they are much better than those without pre-training.

方法

为了利用预训练方法MambaMIM，一个叫做HyMamba的CNN-Mamba混合模型(CNN在顶部，视觉曼巴在底部)被选为基础线骨干。混合模型利用了CNN最后阶段后的Mamba来增强远程依赖性。