当下一代Transformer遇上ViT：Vision RetNet已来！

2023-10-08 05:10

©PaperWeekly 原创 · 作者 | 岳廷

论文标题：

RMT: Retentive Networks Meet Vision Transformers

论文链接：

https://arxiv.org/pdf/2309.11523.pdf

要解决的问题：前不久来自微软的研究人员刚提出了 RetNet，一个被誉为“打破不可能三角形”的模型架构，即高并行训练度，低推理延迟，以及高性能，但该模型主要用于 NLP。详细的 RetNet 介绍可参考下一代Transformer：RetNet结构可视化及Vision RetNet展望。

▲ RetNet 论文中的“不可能三角”

解决方案：来自中科院的研究团队提出 RMT，一种 RetNet 和 Vision Transformer 结合的解决方案，RetNet 引入了显式衰减机制，使得模型能够根据空间距离的先验知识对不同的 token 进行不同程度的注意力。RMT 在视觉主干网络中加入了显式衰减，同时为了降低全局建模的计算代价，将建模过程沿着图像的两个坐标轴进行分解。

实现的效果：大量实验结果证明提出的 RMT 在各种计算机视觉任务中展现出的卓越性能，例如，在 ImageNet-1K 数据集上，RMT 只用 4.5G FLOPs 就达到了 84.1% 的 Top-1 准确率。在所知的模型中，相似大小和训练策略下，RMT 达到了最高的 Top-1 准确率。此外，RMT 在下游任务中如目标检测、实例分割和语义分割中也明显优于现有的视觉骨干网络（特别是在目标检测网络中具有显著优势）。效果示意如下：

额外：单看 RetNet 论文数据，RetNet 潜力很大，RetNet 论文放出 3 个月，中科院已有 Vision Transformer 模型架构放出，足以见该方案潜力很大！相信后续还会有更多 RetNet 尝试方案出来。另外，RMT 论文方案还在进行中，后续可能还会有新的数据持续放出。

方法

1.1 前言

Retentive Network: RetNet 是一个强大的语言模型架构。这项工作为序列建模提出了保留机制。保留为语言模型带来了显式衰减，这是 Transformer 没有的。保留首先以循环方式考虑一个序列建模问题。它可以写成公式（1）：

在训练过程中，对于并行训练过程，公式 1 可以写成公式 2：

其中是的共轭复数，同时包含因果遮挡和指数衰减，它们表示一维序列中 token 之间的相对距离，从而引入了先验知识。基于 1D 的显式衰减，尝试将其发展到 2D 并将其引入视觉模型中。

1.2 保留自注意力

单向到双向：由于语言任务的因果性质，RetNet 中的保留是单向的，这意味着每个 token 只能关注它前面的 tokens，不能关注后面的 tokens。这不适用于没有因果性的任务，如图像识别任务。因此，首先将保留扩展到双向，其中每个 token 的输出变为公式（3）：

其中是 token 的数量。等式可以重排为并行形式，表示为公式（4）：

其中 BiRetention 表示具有双向建模能力的保留。

1D 到 2D：尽管保留现在具有双向建模的能力，但这种建模能力仍局限于一维，仍不适用于二维图像。因此，进一步将一维保留扩展到二维。

对于图像，每个 token 在平面上都有一个唯一的二维坐标。对于第 n 个 token，用表示它的二维坐标。基于每个 token 的二维坐标，修改矩阵中的每个元素为对应 token 对在其各自位置的曼哈顿距离，完成从一维到二维的衰减系数变换。矩阵变换为公式（5）：

在 RetNet 中，消除了 Softmax 并将其替换为门控函数，以增加运算符的非线性。但根据文章的实验，这种方法对视觉模型没有带来更好的结果。相反，它引入了额外的参数和计算复杂度。因此，仍然使用 Softmax 为模型引入非线性。基于上述步骤，保留自注意力可以表示为公式（6）：

早期阶段对 ReSAle 进行分解：当前的 ReSA 不完全适用于图像识别任务。这是因为在视觉骨干网络的早期阶段，存在大量 token，导致 Attention 计算成本过高。这也是大多数 Vision Transformer 变种试图解决的问题。

ReSA 也遇到了这个问题。因此，按图像的两个轴分解 ReSA，如公式（7）所述的具体过程：

基于这种 ReSA 的分解，每个 token 的感受野形状如图 3 所示，与完整 ReSA 的感受野形状相同。

为进一步增强 ReSA 的局部表达能力，还引入了使用 DWConv 的局部增强模块：

1.3 整体架构

整个模型的架构如图 2 所示。类似于传统的骨干网络，它被划分为四个阶段。前三个阶段使用分解的 ReSA，最后一个阶段使用原始 ReSA。与许多以前的骨干网络一样，在模型中加入了 CPE。

实验

在多个视觉任务如图像分类、目标检测、实例分割和语义分割上进行了大量实验，以验证所提出方法的有效性。

2.1 图像分类

设置：从头开始在 ImageNet-1K 上训练模型。遵循 DeiT 中的相同训练策略。最大随机深度增加率分别设置为 RMT-T/S/B/L [21] 的 0.1/0.15/0.4/0.5。使用 AdamW 优化器和余弦衰减学习率调度器来训练模型。将初始学习率、权重衰减和批量大小分别设置为 0.001、0.05 和 1024。

结果：在表 1 中与许多最先进的模型进行了比较。表中的结果表明，RMT 在所有设置下都持续优于以前的模型。具体来说，RMT-S 只用 4.5 GFLOPs 就达到了 84.1% 的 Top-1 准确率。RMT-B 也优于 iFormer，FLOPs 相似而 Top-1 准确率提高了 0.4%。此外，RMT-L 模型在 Top-1 准确率上优于 MaxViT-B0.6%，而 FLOPs 更少。文章的 RMT-T 也优于许多轻量级模型。

2.2 目标检测和实例分割

设置：采用 MMDetection 来实现 RetinaNet 和 Mask-RCNN。对这两种策略使用常用的 “1×”（12 轮训练）设置，并对 Mask-RCNN 使用 “3×+MS”。在训练期间，图像被缩放到短边 800 像素，长边在 1333 像素内。采用 AdamW 优化器，学习率为 0.0001，权重衰减为 0.05，批量大小为 16 来优化模型。对于 “1×” 计划，学习率在第 8 轮和第 11 轮时以 0.1 的衰减率下降。而对于 “3×+MS” 计划，学习率在第 27 轮和第 33 轮时以 0.1 的衰减率下降。

结果：表 2 和表 3 显示了 RetinaNet 和 Mask R-CNN 的结果。结果表明，RMT 在所有比较中都取得了最佳性能。对于 RetinaNet 框架，RMT-T 优于 FAT-B2 +1.1 AP，而 S/B/L 也优于其他方法。对于 “1×” 计划的 Mask R-CNN，RMT-L优于最近的 InternImage-B +1.8 框 AP 和 +1.9 掩膜 AP。对于 “3×+MS” 计划，RMT-S 优于 InternImage-T +1.6 框 AP 和 +1.2 掩模 AP。所有上述结果表明 RMT 以明显的优势超过了对应的对手。

2.3 语义分割

设置：基于 MMSegmentation 采用语义 FPN 和 UperNet，将在 ImageNet-1K 上预训练的 RMT 用作骨干网络。使用与 PVT 相同的设置来训练语义 FPN 模型，并对模型进行 80k 次迭代的训练。所有模型的输入分辨率为 512×512。测试模型时，将图像的短边缩放为 512 像素。对于 UperNet，遵循 Swin 中的默认设置。使用 AdamW 和 0.01 的权重衰减作为优化器，对模型进行 160k 次迭代的训练。学习率设置为 6×10-5，热身 1500 次迭代。

结果：语义分割的结果见表 4 和表 5。所有 FLOPs 都是以 512×2048 的分辨率测量的。所有模型在所有比较中都取得了最佳性能。具体来说，RMT-S 用语义 FPN优于Shunted-S +1.2 mIoU。此外，RMT-B 优于最近的InternImage-S +1.3 mIoU。所有上述结果表明 RMT 模型在密集预测中具有优越性。