RMT: Retentive Networks Meet Vision Transformers——Retentive Networks 与 Vision Transformers 的结合

这篇文章提出了一种名为 RMT（Retentive Networks Meet Vision Transformers） 的新型视觉骨干网络，旨在解决 Vision Transformer (ViT) 中 Self-Attention 机制的两个主要问题：缺乏显式空间先验 和 二次计算复杂度。以下是文章的主要内容总结：

1. 问题背景

• Vision Transformer (ViT) 在计算机视觉任务中表现出色，但其核心组件 Self-Attention 缺乏显式空间先验，且计算复杂度为二次方，限制了其在大规模图像任务中的应用。

• 现有的改进方法（如 Swin Transformer）通过引入窗口机制和相对位置编码来缓解这些问题，但仍存在局限性。

2. 主要贡献

• 曼哈顿自注意力 (MaSA)：受 NLP 领域 Retentive Network (RetNet) 的启发，作者将 RetNet 的时间衰减机制扩展到空间域，提出了基于曼哈顿距离的空间衰减矩阵，为 Self-Attention 引入了显式空间先验。

• 注意力分解：为了在不破坏空间衰减矩阵的情况下降低计算复杂度，作者提出了一种沿图像两个轴分解 Self-Attention 的方法，实现了线性复杂度的全局信息建模。

• RMT 架构：基于 MaSA 和注意力分解，作者构建了 RMT，一个强大的通用视觉骨干网络，能够在保持线性复杂度的同时引入显式空间先验。

3. 实验结果

• 图像分类：RMT 在 ImageNet-1K 上取得了优异的性能，RMT-S 和 RMT-L 分别达到了 84.1% 和 86.1% 的 top-1 准确率，且计算量较低。

• 目标检测与实例分割：在 COCO 数据集上，RMT 在 RetinaNet、Mask R-CNN 和 Cascade Mask R-CNN 框架下均取得了领先的性能。

• 语义分割：在 ADE20K 数据集上，RMT 在 Semantic FPN 和 UperNet 框架下均表现优异，超越了现有的 SOTA 模型。

4. 消融实验

• 作者通过消融实验验证了 MaSA、局部上下文增强模块 (LCE)、条件位置编码 (CPE) 和卷积 stem 的有效性，表明这些组件对模型性能的提升起到了关键作用。

5. 结论

• RMT 通过引入显式空间先验和线性复杂度的注意力机制，显著提升了 Vision Transformer 的性能，并在多个视觉任务中取得了优异的成果。

核心创新点：

• 空间衰减矩阵：基于曼哈顿距离的空间衰减矩阵为 Self-Attention 引入了显式空间先验。

• 注意力分解：沿图像两个轴分解 Self-Attention，实现了线性复杂度的全局信息建模。

• RMT 架构：结合 MaSA 和注意力分解，构建了一个高效且强大的视觉骨干网络。

这篇文章通过结合 Retentive Network 的思想和 Vision Transformer 的架构，提出了一种新的视觉骨干网络 RMT，显著提升了 Transformer 在视觉任务中的性能和效率。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

官方项目地址在这里，如下所示：

摘要：

近年来，Vision Transformer (ViT) 在计算机视觉领域获得了越来越多的关注。然而，ViT 的核心组件 Self-Attention 缺乏显式的空间先验，并且具有二次计算复杂度，从而限制了 ViT 的适用性。为了缓解这些问题，我们从 NLP 领域最近的 Retentive Network (RetNet) 中汲取灵感，提出了 RMT，一种具有显式空间先验的通用视觉骨干网络。具体来说，我们将 RetNet 的时间衰减机制扩展到空间域，并提出了基于曼哈顿距离的空间衰减矩阵，以将显式空间先验引入 Self-Attention。此外，我们提出了一种能够适应显式空间先验的注意力分解形式，旨在减少建模全局信息的计算负担，同时不破坏空间衰减矩阵。基于空间衰减矩阵和注意力分解形式，我们可以灵活地将显式空间先验集成到视觉骨干网络中，并保持线性复杂度。大量实验表明，RMT 在各种视觉任务中表现出卓越的性能。具体来说，在没有额外训练数据的情况下，RMT 在 ImageNet-1k 上分别以 27M/4.5GFLOPs 和 96M/18.2GFLOPs 的参数量和计算量实现了 84.8% 和 86.1% 的 top-1 准确率。在下游任务中，RMT 在 COCO 检测任务中实现了 54.5 的 box AP 和 47.2 的 mask AP，在 ADE20K 语义分割任务中实现了 52.8 的 mIoU。

1. 引言

Vision Transformer (ViT) [13] 是一种备受研究人员青睐的优秀视觉架构。然而，作为 ViT 的核心模块，Self-Attention 的固有结构缺乏显式的空间先验。此外，Self-Attention 的二次复杂度在建模全局信息时会导致显著的计算成本。这些问题限制了 ViT 的应用。

许多工作已经尝试缓解这些问题 [14, 18, 33, 39, 53, 60, 64]。例如，在 Swin Transformer [39] 中，作者通过对用于自注意力的 token 进行窗口划分操作来减少自注意力的计算成本，并通过使用窗口和相对位置编码为模型引入空间先验。此外，NAT [21] 改变了 Self-Attention 的感受野以匹配卷积的形状，减少了计算成本，同时通过其感受野的形状使模型能够感知空间先验。

与之前的方法不同，我们从 NLP 领域最近成功的 Retentive Network (RetNet) [50] 中汲取灵感。RetNet 利用距离依赖的时间衰减矩阵为一维单向文本数据提供显式的时间先验。在 RetNet 之前，ALiBi [45] 也应用了类似的方法并在 NLP 任务中取得了成功。我们将这种时间衰减矩阵扩展到空间域，开发了一种基于 token 之间曼哈顿距离的二维双向空间衰减矩阵。在我们的空间衰减矩阵中，对于目标 token，周围的 token 越远，其注意力分数的衰减程度越大。这一特性使目标 token 能够感知全局信息，同时为不同距离的 token 分配不同级别的注意力。我们使用这种空间衰减矩阵将显式空间先验引入视觉骨干网络。我们将这种受 RetNet 启发并结合曼哈顿距离作为显式空间先验的 Self-Attention 机制命名为 曼哈顿自注意力 (MaSA)。

除了显式空间先验外，使用 Self-Attention 进行全局建模的另一个问题是巨大的计算负担。之前的稀疏注意力机制 [12, 39, 56, 66, 78] 和 RetNet [50] 中的保留分解方式大多会破坏空间衰减矩阵，使其不适用于 MaSA。为了在不破坏空间衰减矩阵的情况下稀疏地建模全局信息，我们提出了一种沿图像两个轴分解 Self-Attention 的方法。这种分解方法在不丢失任何先验信息的情况下分解了 Self-Attention 和空间衰减矩阵。分解后的 MaSA 以线性复杂度建模全局信息，并具有与原始 MaSA 相同的感受野形状。我们在图 2 中将 MaSA 与其他 Self-Attention 机制进行了比较。可以看出，我们的 MaSA 为模型引入了比其对应物更丰富的空间先验。

基于 MaSA，我们构建了一个强大的视觉骨干网络 RMT。我们通过大量实验证明了所提出方法的有效性。如图 1 所示，我们的 RMT 在图像分类任务上优于最先进的 (SOTA) 模型。此外，与其他模型相比，我们的模型在目标检测、实例分割和语义分割等任务中表现出更显著的优势。我们的贡献可以总结如下：

• 我们提出了一种基于曼哈顿距离的空间衰减矩阵来增强 Self-Attention，创建了具有显式空间先验的曼哈顿自注意力 (MaSA)。

• 我们提出了一种 MaSA 的分解形式，使得在不破坏空间衰减矩阵的情况下，能够以线性复杂度建模全局信息。

• 利用 MaSA，我们构建了 RMT，一个强大的通用视觉骨干网络。RMT 在没有额外训练数据的情况下，在 ImageNet-1k 图像分类中获得了高 top-1 准确率，并在目标检测、实例分割和语义分割等任务中表现出色。

2. 相关工作

Transformer。Transformer 架构最初在 [55] 中提出，用于解决循环模型的训练限制，随后在许多 NLP 任务中取得了巨大成功。通过将图像分割成小的、不重叠的 patch 序列，Vision Transformer (ViTs) [13] 也引起了广泛关注，并广泛应用于视觉任务 [6, 16, 20, 44, 61, 68]。与过去 RNN 和 CNN 分别主导 NLP 和 CV 领域不同，Transformer 架构在各种模态和领域中大放异彩 [29, 41, 46, 63]。在计算机视觉领域，许多研究试图将空间先验引入 ViT 以减少训练数据需求 [7, 21, 52]。同时，各种稀疏注意力机制被提出以减少 Self-Attention 的计算成本 [14, 56, 57, 60]。

Transformer 中的先验知识。许多尝试将先验知识引入 Transformer 模型以提高其性能。原始 Transformer [13, 55] 使用三角函数位置编码为每个 token 提供位置信息。在视觉任务中，[39] 提出使用相对位置编码代替原始的绝对位置编码。[7] 指出，卷积层中的零填充也可以为 ViT 提供位置感知，这种位置编码方法非常高效。在许多研究中，FFN 中的卷积 [14, 18, 57] 被用于视觉模型，以进一步丰富 ViT 中的位置信息。对于 NLP 任务，在最近的 Retentive Network [50] 中，时间衰减矩阵被引入，为模型提供基于距离变化的先验知识。在 RetNet 之前，ALiBi [45] 也使用了类似的时间衰减矩阵。

3. 方法论

3.1 预备知识

RetNet 中的时间衰减。Retentive Network (RetNet) 是一种强大的语言模型架构。该工作提出了用于序列建模的保留机制。保留为语言模型带来了时间衰减，这是 Transformer 所不具备的。保留首先以循环方式考虑序列建模问题。它可以写成公式 1：

3.2 曼哈顿自注意力

从 RetNet 的保留机制出发，我们将其演变为曼哈顿自注意力 (MaSA)。在 MaSA 中，我们将保留中的单向一维时间衰减转换为双向二维空间衰减。这种空间衰减将基于曼哈顿距离的显式空间先验引入视觉骨干网络。此外，我们设计了一种简单的方法，同时沿图像的两个轴分解 Self-Attention 和空间衰减矩阵。

从单向到双向衰减：在 RetNet 中，保留是单向的，因为文本数据具有因果性，每个 token 只能关注前面的 token，而不能关注后面的 token。这一特性不适合缺乏因果属性的任务，如图像识别。因此，我们首先将保留扩展为双向形式，表示为公式 3：

在保留中，Softmax 被放弃，取而代之的是一个门控函数。这种变化使 RetNet 具有多种灵活的计算形式，使其能够适应并行训练和循环推理过程。尽管具有这种灵活性，但在我们的实验中仅使用 RetNet 的并行计算形式时，保留门控函数的必要性变得值得商榷。我们的发现表明，这种修改不会改善视觉模型的结果，反而会引入额外的参数和计算复杂性。因此，我们继续使用 Softmax 为模型引入非线性。结合上述步骤，我们的曼哈顿自注意力表示为：

基于 MaSA 的分解，每个 token 的感受野形状如图 3 所示，与完整 MaSA 的感受野形状相同。图 3 表明我们的分解方法完全保留了显式空间先验。

为了进一步增强 MaSA 的局部表达能力，我们借鉴 [78]，引入了一个使用 DWConv 的局部上下文增强模块：

3.3 整体架构

我们基于 MaSA 构建了 RMT，其架构如图 3 所示。与之前的通用视觉骨干网络 [39, 56, 57, 73] 类似，RMT 分为四个阶段。前三个阶段使用分解的 MaSA，最后一个阶段使用原始的 MaSA。与许多之前的骨干网络 [18, 33, 74, 78] 一样，我们将 CPE [7] 集成到我们的模型中。

4. 实验

我们在多个视觉任务上进行了广泛的实验，例如 ImageNet-1K [10] 上的图像分类、COCO 2017 [37] 上的目标检测和实例分割以及 ADE20K [76] 上的语义分割。我们还进行了消融研究，以验证 RMT 中每个组件的重要性。更多细节可以在附录中找到。

4.1 图像分类

设置。我们在 ImageNet-1K [10] 上从头开始训练我们的模型。我们遵循 [52] 中的相同训练策略，唯一的监督是分类损失，以确保公平比较。RMT-T/S/B/L [26] 的最大随机深度增加率分别设置为 0.1/0.15/0.4/0.5。我们使用 AdamW 优化器和余弦衰减学习率调度器来训练模型。我们将初始学习率、权重衰减和批量大小分别设置为 0.001、0.05 和 1024。我们采用了 [39] 中使用的强大数据增强和正则化方法。我们的设置包括 RandAugment [9] (randm9-mstd0.5-inc1)、Mixup [72] (prob=0.8)、CutMix [71] (prob=1.0)、Random Erasing [75] (prob=0.25)。除了传统的训练方法外，类似于 LV-ViT [30] 和 VOLO [70]，我们训练了一个使用 token 标签提供额外监督的模型。

结果。我们在表 1 中将 RMT 与许多最先进的模型进行了比较。表中的结果表明，RMT 在所有设置中始终优于之前的模型。具体来说，RMT-S 仅使用 4.5 GFLOPs 就实现了 84.1% 的 Top1 准确率。RMT-B 在相似的 FLOPs 下也超过了 iFormer [49] 的 0.4%。此外，我们的 RMT-L 模型在 top1 准确率上超过了 MaxViT-B [54] 的 0.6%，同时使用了更少的 FLOPs。我们的 RMT-T 也优于许多轻量级模型。至于使用 token 标签训练的模型，我们的 RMT-S 比当前最先进的 BiFormer-S 高出 0.5%。

4.2 目标检测和实例分割

设置。我们采用 MMDetection [5] 来实现 RetinaNet [36]、Mask-RCNN [24] 和 Cascade Mask R-CNN [2]。我们使用常用的“1×”（12 个训练周期）设置来训练 RetinaNet 和 Mask R-CNN。此外，我们使用“3 × +MS”来训练 Mask R-CNN 和 Cascade Mask R-CNN。遵循 [39]，在训练期间，图像被调整为短边 800 像素，长边在 1333 像素以内。我们使用 AdamW 优化器，学习率为 0.0001，批量大小为 16 来优化模型。对于“1×”调度，学习率在第 8 和第 11 个周期以 0.1 的衰减率下降。而对于“3 × +MS”调度，学习率在第 27 和第 33 个周期以 0.1 的衰减率下降。

结果。表 2、表 3 和表 4 显示了使用不同检测框架的结果。结果表明，我们的 RMT 在所有比较中表现最佳。对于 RetinaNet 框架，我们的 RMT-T 比 MPViT-XS 高出 +1.3 AP，而 S/B/L 也比其他方法表现更好。至于使用“1×”调度的 Mask R-CNN，RMT-L 比最近的 InternImage-B 高出 +2.8 box AP 和 +1.9 mask AP。对于“3 × +MS”调度，RMT-S 比 InternImage-T 高出 +1.6 box AP 和 +1.2 mask AP。此外，对于 Cascade Mask R-CNN，我们的 RMT 仍然比其他骨干网络表现更好。所有这些结果表明，RMT 明显优于其对应物。

4.3 语义分割

设置。我们采用基于 MMSegmentation [8] 的 Semantic FPN [31] 和 UperNet [62]，并将预训练在 ImageNet-1K 上的 RMT 作为骨干网络。我们使用 PVT [56] 的相同设置来训练 Semantic FPN，并训练模型 80k 次迭代。所有模型都以 512 × 512 的输入分辨率进行训练。在测试模型时，我们将图像的短边调整为 512 像素。至于 UperNet，我们遵循 Swin [39] 中的默认设置。我们使用 AdamW 优化器，权重衰减为 0.01，训练模型 160K 次迭代。学习率设置为 6×10−5，预热 1500 次迭代。

结果。语义分割的结果可以在表 5 中找到。所有 FLOPs 都是在 512 × 2048 的分辨率下测量的，除了 RMT-T 组，其 FLOPs 是在 512 × 512 的分辨率下测量的。我们的所有模型在所有比较中都取得了最佳性能。具体来说，我们的 RMT-S 在使用 Semantic FPN 时比 Shunted-S 高出 +1.2 mIoU。此外，我们的 RMT-B 比最近的 InternImage-S 高出 +1.8 mIoU。所有这些结果证明了我们的模型在密集预测中的优越性。

4.4 消融实验

RMT 的鲁棒性。我们在 ImageNet-A (IN-A) 和 ImageNet-R (IN-R) 上评估了 RMT（在 ImageNet-1k 上训练）。RMT 的表现优于 Swin 和以鲁棒性为导向的 FAN（见表 6）。

ImageNet-V2 上的结果。RMT 在 ImageNet-V2 上的表现如表 7 所示。“IN-V2” 表示 ImageNet-V2。我们的 RMT 在 ImageNet-V2 上表现出非常令人印象深刻的性能。

与之前工作的严格对比。为了与之前的方法进行严格对比，我们将 RMT 的超参数（例如是否使用分层结构、分层模型中四个阶段的通道数、是否使用位置编码和卷积 stem 等）与 DeiT [52] 和 Swin [39] 对齐，并仅将 Self-Attention/Window Self-Attention 替换为我们的 MaSA。对比结果如表 8 所示，RMT 显著优于 DeiT-S、Swin-T 和 Swin-S。

MaSA。我们验证了曼哈顿自注意力 (MaSA) 对模型的影响，如表 8 所示。MaSA 显著提升了模型在图像分类和下游任务中的性能。具体来说，MaSA 的分类准确率比普通注意力高出 0.8%。

LCE。局部上下文增强 (LCE) 也在我们模型的优异表现中发挥了作用。LCE 将 RMT 的分类准确率提高了 0.3%，并增强了模型在下游任务中的表现。

CPE。与之前的方法一样，条件位置编码 (CPE) 为我们的模型提供了灵活的位置编码和更多的位置信息，有助于提升模型在图像分类和下游任务中的表现。

卷积 Stem。模型的初始卷积 stem 提供了更好的局部信息，从而进一步增强了模型在各种任务中的表现。

分解的 MaSA。在 RMT-S 中，我们将第三阶段的分解 MaSA (MaSA-d) 替换为原始的 MaSA，以验证我们分解方法的有效性，如表 9 所示。在图像分类方面，MaSA-d 和 MaSA 实现了相当的准确率。然而，对于语义分割，使用 MaSA-d 显著降低了计算负担，同时产生了相似的结果。

为什么 MaSA 有效？ 如表 10 所示，逐步将 Swin-T 转换为 RMT-Swin-T 揭示了关键见解：1) 分解形式。将 Swin 的 WSA 替换为分解的 Self-Attention 提高了性能，可能是由于分解形式的全局感知能力。2) 衰减矩阵引入了局部偏差。UniFormer 强调了早期模型阶段局部偏差的重要性。通过适当的衰减率引入局部偏差是有益的。3) 多重衰减率。为每个 MaSA 头分配不同的衰减率引入了多尺度局部偏差，进一步提升了性能。

推理速度。我们在表 11 中将 RMT 的推理速度与最近表现最佳的视觉骨干网络进行了比较。我们的 RMT 展示了速度与准确性之间的最佳权衡。

5. 结论

在这项工作中，我们提出了 RMT，一种具有显式空间先验的视觉骨干网络。RMT 将用于因果建模的时间衰减扩展到空间层面，并引入了基于曼哈顿距离的空间衰减矩阵。该矩阵将显式空间先验引入 Self-Attention。此外，RMT 使用了一种 Self-Attention 分解形式，可以在不破坏空间衰减矩阵的情况下稀疏地建模全局信息。空间衰减矩阵和注意力分解形式的结合使 RMT 具有显式空间先验和线性复杂度。在图像分类、目标检测、实例分割和语义分割中的大量实验验证了 RMT 的优越性。

总结

这篇论文提出了一种新的视觉骨干网络 RMT，通过引入基于曼哈顿距离的空间衰减矩阵和注意力分解机制，显著提升了 Vision Transformer 的性能，并在多个视觉任务中取得了优异的成果。