CVPR | 2025 | MAP：通过掩码自回归预训练释放混合 Mamba - Transformer 视觉骨干网络的潜力

• 论文：https://arxiv.org/pdf/2410.00871
• 代码：https://github.com/yunzeliu/MAP
• （代码）镜像：https://gitee.com/apuppyliu-cong/MAP.git
• 会议：CVPR
• 年份：2025

创新点

• 贡献主要有以下三点：
  • 第一，首次提出了一种用于预训练混合 Mamba - Transformer 视觉骨干网络的新方法，在统一框架下提升了混合骨干网络以及纯 Mamba 和纯 Transformer 骨干网络的性能。
  • 第二，为了便于分析，我们还对纯 Mamba 在自回归预训练中的关键组件进行了深入分析，发现其有效性取决于预训练顺序与 Mamba 扫描顺序的一致性，以及合适的令牌掩码比例。
  • 第三，我们证明了所提出的 MAP 方法在各种 2D 和 3D 数据集上，显著提升了 Mamba - Transformer 和纯 Mamba 骨干网络的性能。大量的消融实验证实了每个设计选择的合理性和有效性。

初步研究

表 1. 初步研究。我们使用 ViT - B 和 Vim - B 作为默认配置。AR 策略按行优先顺序处理图像令牌，而 MAE 按照默认设置运行。对于对比学习，我们仅使用裁剪和缩放的数据增强，并使用 MoCov2 进行预训练。所有实验都在 224x224 的分辨率下进行。AR 的掩码令牌数量设置为 40 个令牌（20%）。

1. 实验表明，MAE 更适合 Transformer 的预训练，而 AR 更适合 Mamba 的预训练。
   1.1 对于 ViT 而言，应用 MAE 策略对于建立令牌之间的双向关联、从而提高性能至关重要
   1.2 对于 Vim，更重要的是对前后令牌之间的连续性进行建模
2. 经过深入分析发现，与扫描顺序一致的自回归预训练以及适当的掩码比例是 Mamba 预训练的关键
   

表 2. AR 顺序对下游任务的影响。Vim（R）指采用行优先扫描的 Vim。Vim（C）指采用列优先扫描的 Vim。AR（R）指行优先自回归预训练。AR（C）指列优先自回归预训练。结果表明，当 AR 预训练设计与 Mamba 的扫描顺序一致时，可获得最佳性能。

图 3. AR 预训练和 Mamba 扫描的不同顺序。行优先和列优先的顺序使网络能够以不同的方式和序列感知局部信息。

3. 自回归预训练的掩码比例。
   

表 3. 掩码比例对 AR 预训练的影响。分别掩码 1 个令牌（0.5%）、20 个令牌（10%）、40 个令牌（20%）、60 个令牌（30%）、100 个令牌（50%）和 140 个令牌（70%），同时记录在下游任务上的微调结果。实验表明，适当的掩码比例对于自回归预训练很重要。

初步结论

得出以下三个结论，作为设计混合骨干网络预训练的参考：

• MAE 更适合 Transformer，而 AR 更适合 Mamba。
• 对于 Transformer 的 MAE 预训练，非对称结构和适当的掩码比例很重要。
• 对于 Mamba 的 AR 预训练，适当的 AR 顺序和掩码比例很重要。

方法

本文的重点是研究如何预训练混合 Mamba - Transformer 框架，而非设计混合框架本身。

确定一个混合网络

结果表明，采用 MMMTMMMT 的混合方式效果最佳。在比较 Mamba - R * 与 MMMMMMTT 时，发现在 Mamba 之后添加 Transformer 增强了其长上下文建模能力，从而提升了性能。在比较 MMMMMMTT 与 TTMMMMMM 时，我们观察到仅仅在 Mamba 之后附加 Transformer 并不能充分发挥该架构的潜力。

这表明在开始部分融入 Transformer 对于提取足够的局部特征至关重要。我们认为 MMMTMMMT 方法有效地平衡了局部特征提取和上下文建模增强，因此将其作为我们的默认配置。

图 4. 不同的混合模型设计。（d）取得了最佳结果，被设置为默认模型，并称为 HybridNet。

表 4. Mamba - Transformer 骨干网络的混合设计。所有实验都是从零开始训练的。Mamba - R表示 24 个 Mamba - R [40] Mamba 层加上 8 个额外的 Mamba 层。DeiT表示 24 个 DeiT [39] Transformer 层加上 8 个额外的 Transformer 层。MMMMMMTT 表示 24 个 Mamba 层之后跟随 8 个 Transformer 层。TTMMMMMM 表示 8 个 Transformer 层之后跟随 24 个 Mamba 层。TMMMTMMM 表示由 1 个 Transformer 层和 3 个 Mamba 层组成的单元，重复 8 次。MMMTMMMT 表示由 3 个 Mamba 层之后跟随 1 个 Transformer 层组成的单元，重复 8 次。

方法

图2.（a）MAE预训练。其核心在于基于未掩码的令牌重建被掩码的令牌，以建立全局双向上下文理解。（b）AR预训练。它着重于建立上下文之间的关联，其可扩展性在大型语言模型领域已得到充分验证。（c）MAP预训练（我们提出的方法）。我们的方法首先对输入图像进行随机掩码处理，然后以逐行自回归的方式重建原始图像。这种预训练方法在对局部特征的上下文特征以及局部特征之间的关联进行建模方面展现出显著优势，使其与Mamba-Transformer混合架构具有高度的兼容性。（d）在ImageNet-1K上不同预训练策略下的性能提升。我们发现MAE预训练更适合Transformer，而AR更适配Mamba。另一方面，MAP更适合Mamba-Transformer骨干网络。此外，MAP在对纯Mamba或纯Transformer骨干网络进行预训练时也表现出令人瞩目的性能，这体现了我们方法的有效性和广泛适用性。

如图 2（c）所示，对于给定图像，HybridNet 首先将经过随机掩码处理的图像映射到特征空间，随后借助 Transformer 解码器按行对原始图像进行解码。

假设将图像 $\mathbf{I}$ 按行划分：
I={r1,r2,...,rM}I=\left\{r_{1}, r_{2}, ..., r_{M}\right\}I={r1​,r2​,...,rM​}
每行 $r_i$ 由以下令牌组成：
ri={xi1,xi2,...,xiN}r_{i}=\left\{x_{i 1}, x_{i 2}, ..., x_{i N}\right\}ri​={xi1​,xi2​,...,xiN​}
在每行中选取一部分令牌进行掩码处理。

令 Mi⊂{1,2,...,N}M_{i} \subset\{1,2, ..., N\}Mi​⊂{1,2,...,N} 表示行 $r_i$ 中被掩码令牌的索引。

对于给定的行 $r_i$ ，需同时预测所有被掩码的令牌：
$$p\left(x_{ij}|x_{i,j\notin M_i},r_{<i}\right)$$

其中，$r_{<i}$ 指的是第 $i$ 行之前的所有行。

第 $i$ 行中令牌的预测取决于所有先前的行以及该行内可见的令牌。这可以表示为：
$$p\left(r_i|r_{<i}\right)=\prod _{j=1}^{N}p\left(x_{ij}|x_{i,j\notin M_i},r_{<i}\right)$$
整体的损失函数为预测令牌的负对数似然之和：
$$\mathcal{L}=-\sum _{i=1}^M\sum _{j\in M_i}\log p\left(x_{ij}|x_{i,j\notin M_i},r_{<i}\right)$$

掩码机制

实验了不同的掩码策略，包括随机掩码、顺序掩码和对角线掩码。实验表明，随机掩码能产生最佳结果。

图 5. 不同的掩码策略。随机掩码策略产生最佳结果。

掩码比例

MAP的transformer解码器

使用 Transformer 解码器而非 Mamba 解码器的原因是，Transformer 解码器可以通过应用解码器掩码，基于编码器的特征进行区域重建。相比之下，Mamba 解码器由于其单向扫描特性，难以同时重建整个局部区域。

解码器采用独特的行解码策略，允许一次对一行令牌进行自回归解码，增强了网络捕获局部特征和区域间上下文关系的能力。

图 6. 不同的解码器掩码。绿色表示激活。白色表示非激活。

表 6. 解码器掩码设计。我们的 MAP 解码器策略取得最佳结果。

重建目标

重建目标与 MAE 一致，我们将归一化的原始像素作为重建目标，并采用 MSE 损失。

实验

ImageNet-1k 上的 2D 分类

预训练阶段，我们采用 50% 掩码比例的随机掩码策略，仅使用随机裁剪作为数据增强策略。在所有设置下训练 1600 个 epoch。在微调阶段，我们直接微调 400 个 epoch 并报告结果。

表 8. ImageNet-1k 分类结果。吞吐量在 A100 GPU 上计算。内存开销以 128 的批处理大小衡量。我们的结果以蓝色突出显示。