PaperlayerView exact revision on arXiv

Source-linked AI summary

RMT: Retentive Networks Meet Vision Transformers

Qihang Fan, Huaibo Huang, Mingrui Chen, Hongmin Liu, Ran He

arXiv:2309.11523v6cs.CV

TL;DR

ViT Self-Attention lacks explicit spatial priors and has quadratic complexity, limiting its applicability. RMT addresses these issues with Manhattan-based spatial decay and a compatible attention decomposition, and the paper reports strong results across classification, detection, instance segmentation, and segmentation.

  • Problem

    ViT Self-Attention lacks explicit spatial priors and has quadratic computational complexity, constraining ViT applicability.

  • Method

    RMT extends RetNet’s temporal decay to bidirectional two-dimensional Manhattan-distance spatial decay and decomposes Self-Attention with the decay matrix for linear-complexity global modeling.

  • Results

    RMT exhibits strong performance across image classification, object detection, instance segmentation, and semantic segmentation.

  • Takeaways & Limitations

    Combining explicit spatial priors with linear-complexity global modeling yields a vision backbone applicable across multiple vision tasks.

  • Takeaways & Limitations

    Replacing Softmax with RetNet’s gating function does not train stably in the decomposed MaSA form, though the authors do not conclude the gating function is generally inferior.

Abstract

from arXiv · show

Vision Transformer (ViT) has gained increasing attention in the computer vision community in recent years. However, the core component of ViT, Self-Attention, lacks explicit spatial priors and bears a quadratic computational complexity, thereby constraining the applicability of ViT. To alleviate these issues, we draw inspiration from the recent Retentive Network (RetNet) in the field of NLP, and propose RMT, a strong vision backbone with explicit spatial prior for general purposes. Specifically, we extend the RetNet's temporal decay mechanism to the spatial domain, and propose a spatial decay matrix based on the Manhattan distance to introduce the explicit spatial prior to Self-Attention. Additionally, an attention decomposition form that adeptly adapts to explicit spatial prior is proposed, aiming to reduce the computational burden of modeling global information without disrupting the spatial decay matrix. Based on the spatial decay matrix and the attention decomposition form, we can flexibly integrate explicit spatial prior into the vision backbone with linear complexity. Extensive experiments demonstrate that RMT exhibits exceptional performance across various vision tasks. Specifically, without extra training data, RMT achieves **84.8%** and **86.1%** top-1 acc on ImageNet-1k with **27M/4.5GFLOPs** and **96M/18.2GFLOPs**. For downstream tasks, RMT achieves **54.5** box AP and **47.2** mask AP on the COCO detection task, and **52.8** mIoU on the ADE20K semantic segmentation task. Code is available at https://github.com/qhfan/RMT

1. Introduction

The paper identifies missing spatial priors and quadratic complexity as key limitations of ViT Self-Attention, then introduces MaSA and its decomposition to address both issues. Using MaSA, RMT is presented as a general-purpose vision backbone with strong performance across vision tasks.

  • ViT Self-Attention lacks explicit spatial priors and has quadratic complexity, limiting ViT applications.
  • MaSA extends RetNet’s temporal decay into a bidirectional two-dimensional spatial decay matrix based on Manhattan distance.The decay increases with token distance, allowing global information while assigning distance-dependent attention.
  • The proposed attention decomposition preserves the spatial decay matrix while modeling global information with linear complexity.It decomposes Self-Attention and the spatial decay matrix along both image axes without losing prior information.
  • RMT uses MaSA to provide a general-purpose vision backbone with explicit spatial priors.The paper reports applications in image classification, object detection, instance segmentation, and semantic segmentation.

2. Related Work

Prior work incorporates positional or spatial information into Transformer models through positional encodings, convolutional components, and RetNet’s temporal decay. Vision Transformers have become widely used across computer vision and other modalities.

  • Transformer architectures became successful in NLP and later gained broad use in vision through patch-based Vision Transformers.
  • Transformer models incorporate prior knowledge through trigonometric or relative positional encoding and convolution-based positional awareness.
  • RetNet introduces a temporal decay matrix to provide temporal information in language modeling.

3. Methodology

RMT extends RetNet’s temporal decay into bidirectional two-dimensional spatial decay based on Manhattan distance, then decomposes attention along image axes to preserve this prior while reducing global-modeling cost.

  • From RetNet to MaSA: RetNet provides explicit temporal priors through distance-dependent decay, motivating RMT’s spatial extension for vision.The original mechanism models one-dimensional, unidirectional text sequences.
  • From RetNet to MaSA: MaSA converts unidirectional one-dimensional retention into bidirectional two-dimensional spatial decay for image tokens.This addresses the noncausal and two-dimensional structure of images.
  • Spatial prior: Manhattan-distance decay assigns greater attention-score decay to farther tokens while retaining access to global information.Each token pair’s decay is computed from their two-dimensional coordinates.
  • Design choices: RMT keeps Softmax instead of RetNet’s gating function because gating adds parameters and computation without improving vision-model results.The choice applies when using RetNet’s parallel computation form in the experiments.
  • Decomposed MaSA: The decomposed MaSA computes horizontal and vertical attention separately and applies one-dimensional bidirectional decay matrices along both axes.This decomposition targets the high cost of global attention in early backbone stages.
  • Decomposed MaSA: The decomposition preserves the complete MaSA receptive-field shape and therefore retains the explicit spatial prior.The paper states that the decomposed and complete MaSA receptive fields are identical.
  • RMT backbone: RMT uses decomposed MaSA in its first three stages, original MaSA in the last stage, and incorporates CPE.The backbone is organized into four stages.

4. Experiments

RMT is evaluated across ImageNet-1K classification, COCO detection and instance segmentation, and ADE20K semantic segmentation, with ablations isolating its components. The experiments report strong accuracy, dense-prediction performance, and an inference-speed/accuracy trade-off.

  • Image Classification: RMT-S achieves 84.1% Top1-accuracy with 4.5 GFLOPs on ImageNet-1K classification.
  • Image Classification: RMT consistently outperforms previous models across the reported ImageNet-1K classification settings, including comparisons against iFormer, MaxViT-B, and BiFormer-S.RMT-B surpasses iFormer by 0.4%, RMT-L surpasses MaxViT-B by 0.6% with fewer FLOPs, and token-labeled RMT-S surpasses BiFormer-S by 0.5%.
  • Object Detection and Instance Segmentation: RMT performs best in the reported COCO detection and instance-segmentation comparisons across RetinaNet, Mask R-CNN, and Cascade Mask R-CNN.Reported gains include +1.3 AP over MPViT-XS, +2.8 box AP and +1.9 mask AP over InternImage-B, and +1.6 box AP and +1.2 mask AP over InternImage-T.
  • Semantic Segmentation: All reported RMT models achieve the best performance in the ADE20K semantic-segmentation comparisons.RMT-S exceeds Shunted-S by +1.2 mIoU with Semantic FPN, while RMT-B exceeds InternImage-S by +1.8 mIoU.
  • Ablation Study: MaSA improves classification accuracy by 0.8% over vanilla attention, while LCE improves RMT classification accuracy by 0.3%.
  • Ablation Study: Replacing MaSA with original retention causes lower vision-backbone performance and lower inference speed, whereas RMT provides the optimal reported speed-accuracy trade-off.

5. Conclusion

RMT combines a Manhattan-distance spatial decay matrix with a decomposition form that preserves explicit spatial priors while enabling linear-complexity global modeling. Experiments across classification, detection, instance segmentation, and semantic segmentation validate RMT's performance.

  • 5. Conclusion: RMT extends RetNet’s temporal decay to a two-dimensional Manhattan-distance spatial decay matrix for Self-Attention.The matrix assigns greater attention decay to surrounding tokens that are farther from a target token.
  • 5. Conclusion: RMT decomposes Self-Attention to sparsely model global information without disrupting the spatial decay matrix.The combined design provides explicit spatial priors and linear complexity.
  • 5. Conclusion: Extensive experiments in image classification, object detection, instance segmentation, and semantic segmentation validate RMT’s superiority.

A. Architecture Details

RMT architectures use convolutional stems and stage transitions to embed and downsample image features, with depth-wise convolutions providing positional and local-enhancement components.

  • A. Architecture Details: The convolution stem applies five 3 × 3 convolutions to embed images into 56 × 56 tokens.GELU and batch normalization follow each convolution except the last, which is followed only by batch normalization.
  • A. Architecture Details: Stride-2 3 × 3 convolutions between stages reduce feature-map resolution.
  • A. Architecture Details: RMT uses 3 × 3 depth-wise convolutions in CPE and 5 × 5 depth-wise convolutions in LCE.

B. Experimental Settings

The experiments train ImageNet models from scratch under a DeiT-style protocol and evaluate downstream detection, instance segmentation, and semantic segmentation with established frameworks and settings.

  • B. Experimental Settings: ImageNet models are trained from scratch for 300 epochs using AdamW, cosine learning-rate decay, warm-up, and classification-loss supervision.A second model version additionally uses token labeling supervision.
  • B. Experimental Settings: COCO experiments use RetinaNet, Mask-RCNN, and Cascaded Mask-CNN under 1× and 3×+MS training settings.
  • B. Experimental Settings: ADE20K experiments use UperNet for 160k iterations or SemanticFPN for 80k iterations, both with 512 × 512 inputs.

C. Efficiency Comparison

RMT achieves the best trade-off between inference speed and accuracy among the compared backbones.

  • C. Efficiency Comparison: RMT achieves the best trade-off between speed and accuracy among many competing backbones.

D. Details of Explicit Decay

ReSA assigns different decay parameters γ across heads to control receptive fields and capture multi-scale information. The appendix provides stage- and backbone-specific decay settings alongside architecture and inference-speed details.

  • γ decay across heads: Different γ values across ReSA heads control each head’s receptive field, enabling multi-scale perception.All head γ values are kept within a specified positive interval [a, b].
  • γ decay across heads: For a ReSA module with N heads and receptive-field interval [a, b], the ith head’s γ is specified by Eq. 8.
  • Stage-specific settings: Different backbone stages use different values of a and b, with the settings listed in Table 12.
  • Model details: Table 10 reports the detailed architectures of the models.
  • Model details: Table 11 compares inference speed across the models.
Loading 2309.11523v6…