RMT: Retentive Networks Meet Vision Transformers学习笔记

代码地址：GitHub - qhfan/RMT: (CVPR2024)RMT: Retentive Networks Meet Vision Transformer

论文地址：https://arxiv.org/pdf/2309.11523.pdf

Transformer首次出现在自然语言处理领域，后来迁移到计算机视觉领域，在视觉任务中表现出出色的性能。然而，最近，Retentive Network（RetNet）作为一种有可能取代Transformer的架构出现，引起了自然语言处理社区的广泛关注。因此，作者提出了一个问题，即将RetNet的思想迁移到视觉领域是否也能为视觉任务带来出色的性能。为了解决这个问题，作者将RetNet和Transformer结合起来，提出了RMT。受RetNet启发，RMT在视觉Backbone中引入了显式衰减，并提出了一个基于曼哈顿距离的空间衰减矩阵，将与空间距离相关的先验知识引入到视觉模型中。这种与距离相关的空间先验允许显式控制每个Token可以关注的Token范围。此外，为了降低全局建模的计算成本，作者沿图像的两个坐标轴分解了这个建模过程。

大量的实验表明，RMT在各种计算机视觉任务中表现出色。例如，RMT仅使用4.5G FLOPs在ImageNet-1k上实现了84.1%的Top1准确率。据作者所知，在所有模型中，当模型的大小相似并采用相同的策略进行训练时，RMT在Top1准确率方面表现最佳。此外，RMT在目标检测、实例分割和语义分割等下游任务中明显优于现有的视觉Backbone模型。

1. Introduction

自从Transformer在自然语言处理领域提出以来，它在许多下游任务中取得了出色的表现。尽管计算机视觉和自然语言处理之间存在模态差异，研究人员成功地将这种架构迁移到了视觉任务中，再次给作者带来了与以前的自然语言处理任务一样的巨大惊喜。

最近，自然语言处理领域出现了一些强大的架构，然而还没有任何工作尝试将这些自然语言处理架构转移到视觉任务中。与Transformer相比，RetNet在多个自然语言处理任务中表现出更强的性能。因此，作者希望能够将这个强大的自然语言处理架构RetNet转移到视觉领域。

RetNet中的基本运算符是Retention。Retention和Transformer中的基本的自注意力运算符之间的一个重要区别是引入了衰减系数，这些系数明确地控制每个Token相对于其相邻Token的注意力权重，确保注意力权重随着Token之间的距离增加而衰减。这种衰减有效地将一维距离的先验知识引入到模型中，从而提高了性能。

基于RetNet的研究结果，作者尝试进一步改进Retention机制，使其适应二维形式，并将其引入到视觉任务中。具体来说，原始版本的Retention机制经历了单向衰减，适用于自然语言处理的因果属性。然而，对于没有因果属性的图像，这种单向衰减是不适用的。

因此，作者首先将Retention机制从单向扩展到双向。此外，原始版本的Retention机制设计用于一维顺序信息，不适合在二维空间中使用。因此，考虑到二维空间的空间特性，作者设计了基于二维距离的衰减矩阵。最后，为了解决视觉Backbone早期阶段大量Token带来的高计算负荷问题，作者将二维计算过程分解为沿图像的两个坐标轴分别进行。作者将适用于图像的这种机制命名为 Manhattan Self-Attention（MaSA）机制。基于MaSA机制，作者构建了RMT系列。在图 2 中将 MaSA 与其他自注意力机制进行了比较。可以看出，MaSA为模型引入了比其参照注意力更丰富的空间先验。

作者通过广泛的实验来证明所提出的方法的有效性。如图1所示，作者的RMT在图像分类任务上明显优于最先进的模型（SOTA）。此外，与其他模型相比，作者的模型在目标检测和实例分割等任务中具有更明显的优势。作者的贡献可以总结如下：

作者将Retentive Network的核心机制Retention扩展到了二维情景，引入了与距离相关的空间先验知识到视觉模型中。并提出了一个基于曼哈顿距离的空间衰减矩阵，新的机制被称为Manhattan Self-Attention（MaSA）。
作者将MaSA沿着两个图像轴进行分解，降低了计算复杂性。这种分解方法有效地减小了计算负担，同时对模型性能的影响最小。
广泛的实验证明了RMT的出色性能。特别是在目标检测和实例分割等下游任务中，RMT表现出明显的优势。

2. Related Work

Transformer

Transformer的提出旨在解决递归模型的训练限制，并在许多自然语言处理任务中取得了巨大成功。通过将图像分割为小的、不重叠的Patch序列，视觉Transformer（ViTs）也引起了广泛关注，并在视觉任务中得到了广泛应用。与过去不同，RNN和CNN分别在自然语言处理和计算机视觉领域占主导地位，而Transformer架构在各种模态和领域中都表现出色。

Prior Knowledge in Transformer.

为了增强Transformer模型的性能，已经进行了大量尝试，将先验知识引入其中。最初的Transformer使用三角函数位置编码为每个Token提供位置信息。Swin Transformer提出了使用相对位置编码作为原始的绝对位置编码的替代方法。[Conditional positional encodings]指出卷积层中的零填充也可以为Transformer提供位置感知能力，而这种位置编码方法非常高效。在许多研究中，FFN中的Convolution已被用来进一步丰富Transformer中的位置信息。此外，在最近的Retentive Network中，引入了显式衰减，为模型提供了基于距离变化的先验知识。与传统的基于Transformer的模型相比，RetNet中提出的Retention使用显式衰减来建模一维距离的先验知识，这是一种重要区别。它包括三种计算范式，即并行、递归和分块递归。在Retention中，它使用一个衰减矩阵乘以一个权重矩阵，以根据距离先验控制每个Token看到其周围Token的比例。作者也尝试将这一思想扩展到二维空间。

3. Methodology

3.1. Preliminary

Temporal decay in RetNet.RetNet是一种强大的语言模型架构。本工作提出了用于序列建模的Retention机制。Retention引入了显式衰减到语言模型中，而Transformer没有这个特性。Retention首先以递归方式考虑序列建模问题。它可以写成如下的等式（式1）：

在训练过程中，对于并行训练过程，等式（1）可以写成如下的等式（2）：

其中 $\overline{\Theta }$ 是 $\Theta$ 的复共轭， $D\in \mathbb{R}^{|x|\times |x|}$ 包含了 causal masking和exponential decay，代表了一维序列中的相对距离，这带来了先验知识。

3.2. Manhattan Self-Attention

基于Retention中的一维显式衰减，作者尝试将其扩展到二维，并将空间先验知识引入视觉模型中。

From Unidirectional to Bidirectional Decay: 由于语言任务的因果性质，RetNet中的Retention是单向的，意味着每个Token只能关注其前面的Token，而不能关注其后的Token。这不适用于没有因果属性的任务，例如图像识别任务。因此，作者首先将Retention扩展到二维，在这种情况下，对于每个Token，其输出变为如下的等式（3）：

其中 BiRetention 表示双向建模。

From One-dimensional to Two-dimensional Decay:尽管Retention现在具有双向建模的能力，但这种建模能力仍然局限于一维水平，仍然不适用于二维图像。因此，作者进一步将一维Retention扩展到二维。

对于图像，每个Token在平面内具有唯一的二维坐标。对于第n个Token，作者使用 $(x_{n},y_{n})$ 来表示其二维坐标。基于每个Token的二维坐标，作者修改矩阵D中的每个元素，使其成为相应位置的Token对之间的曼哈顿距离，完成了从一维到二维衰减系数的转换。矩阵D转化为等式（4）：

在Retention中，放弃了Softmax，并用一个门控函数来增加运算符的非线性性。然而，根据作者的实验，这种方法对于视觉模型并没有产生更好的结果。相反，它引入了额外的参数和计算复杂性。因此，作者仍然使用Softmax来引入非线性到作者的模型中。

基于上述步骤，作者的Manhattan Self-Attention可以表示为等式（5）：

Decomposed Manhattan Self-Attention.目前的MaSA并不完全适用于图像识别任务。这是因为在视觉Backbone网络的早期阶段，存在大量的Token，导致了Attention的计算成本过高。这也是大多数视觉Transformer变种都努力解决的问题。作者的MaSA也遇到了这个问题。因此，作者将MaSA分解为图像的两个轴，具体过程如等式（6）所示：

其中

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

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

3.3. Overall Architecture

作者整个模型的架构如图3所示。与传统的Backbone网络类似，它分为4个阶段。前3个阶段使用了分解的MaSA，而最后一个阶段使用了原始的ReSA。与许多先前的Backbone网络一样，作者在模型中引入了CPE。

4. 实验

4.1. Image Classification

作者在表1中将RMT与许多最先进的模型进行了比较。表中的结果表明，RMT在所有设置下始终优于先前的模型。具体而言，RMT-S仅使用4.5 GFLOPs就实现了84.1%的Top1准确率。RMT-B也在类似的FLOPs下超越了iFormer 0.4%。

此外，作者的RMT-L模型在使用更少的FLOPs的情况下，将Top1准确率超过了MaxViT-B，提高了0.6%。作者的RMT-T模型也在性能上超越了许多轻量级模型。

4.2. Object Detection and Instance Segmentation

表2，表3以及表4显示了使用不同检测框架的结果。结果表明，作者的RMT在所有比较中表现最好。对于RetinaNet框架，作者的RMT-T比MPViT-XS 高出 +1.3 AP，而S/B/L也优于其他方法。

对于具有“1×”schedule的Mask R-CNN，RMT-L超越了最近的InternImage-B +1.8框AP和+1.9maskAP。对于“3× +MS”schedule，RMT-S超越了InternImage-T +1.6框AP和+1.2maskAP。所有以上结果表明，RMT明显优于其同类算法。

4.3. Semantic Segmentation

语义分割的结果可以在表5中找到。所有的FLOPs都是以512×2048的分辨率测量的。作者所有的模型在所有比较中都取得了最佳性能。

4.4. Ablation Study

Strict comparison with previous works.为了与之前的方法进行严格的比较，作者将整体架构的RMT超参数（如是否使用分层结构、分层模型四个阶段的通道数、是否使用位置编码和convolution stem等）与DeiT和Swin对齐，只用MaSA替换Self-Attention/Window Self-Attention。比较结果如表6所示，其中RMT明显优于DeiT-S、Swin-T和Swin-S。

MaSA.作者验证了Manhattan Self-Attention对模型的影响，如表6所示。MaSA大大提高了模型在图像分类和下游任务中的性能。具体而言，MaSA的分类准确率比香草注意力高0.8%。

Softmax.在RetNet中，Softmax被替换为非线性门控函数，以适应其各种计算形式。作者用这个门控功能替换了MaSA中的Softmax。然而，利用门控函数的模型无法进行稳定的训练。值得注意的是，这并不意味着门控功能不如Softmax。门控功能可能与作者的分解形式或空间衰减不兼容。

LCE.局部上下文增强也对模型的出色性能起着重要作用。LCE将RMT的分类准确率提高了0.3%，并增强了模型在下游任务中的性能。

CPE.CPE为模型提供了灵活的位置编码和更多的位置信息，有助于提高模型在图像分类和下游任务中的性能。

Convolutional Stem.模型的初始convolutional stem提供了更好的局部信息，从而进一步提高了模型在各种任务上的性能。

DecomposedMaSA.在RMT-S中，作者用原始的MaSA代替了第三阶段的分解MaSA（MaSA-d），以验证分解方法的有效性，如表7所示。在图像分类方面，MaSA-d和MaSA的精度相当。然而，对于语义分割，使用 MaSA-d 可显著减轻计算负担，同时产生相似的结果。

MaSAv.s. Retention.如表 8 所示，作者在 RMT-S 架构中用原始保留替换了 MaSA。使用 Swin-Transformer中采用的方法将tokens划分为块，以实现 chunk-wise retention。由于一维 causal data建模的retention限制，基于它的视觉主干的性能落后于RMT。此外，chunk-wise 和recurrent 形式的retention会破坏视觉主干的并行性，导致推理速度降低。