ViR: Towards Efficient Vision Retention Backbones

TL;DR

ViR tackles the inefficiency of self-attention in Vision Transformers by introducing Vision Retention Networks that realize dual parallel and recurrent retention with a mask-based mechanism. The approach includes 1D and 2D retention, with isotropic and CNN-enhanced Hybrid variants, and a chunkwise formulation that yields fast inference while preserving training parallelism. Empirical results on ImageNet and downstream tasks show competitive accuracy and notable throughput/memory benefits, especially at high resolutions, illustrating ViR's potential as a scalable vision backbone. The work provides a flexible framework that decouples memory usage from sequence length, enabling efficient processing of very large images and long-range dependencies with practical applicability across vision tasks, supported by theoretical equivalence between recurrent and parallel formulations and by comprehensive ablations.

Abstract

Vision Transformers (ViTs) have attracted a lot of popularity in recent years, due to their exceptional capabilities in modeling long-range spatial dependencies and scalability for large scale training. Although the training parallelism of self-attention mechanism plays an important role in retaining great performance, its quadratic complexity baffles the application of ViTs in many scenarios which demand fast inference. This effect is even more pronounced in applications in which autoregressive modeling of input features is required. In Natural Language Processing (NLP), a new stream of efforts has proposed parallelizable models with recurrent formulation that allows for efficient inference in generative applications. Inspired by this trend, we propose a new class of computer vision models, dubbed Vision Retention Networks (ViR), with dual parallel and recurrent formulations, which strike an optimal balance between fast inference and parallel training with competitive performance. In particular, ViR scales favorably for image throughput and memory consumption in tasks that require higher-resolution images due to its flexible formulation in processing large sequence lengths. The ViR is the first attempt to realize dual parallel and recurrent equivalency in a general vision backbone for recognition tasks. We have validated the effectiveness of ViR through extensive experiments with different dataset sizes and various image resolutions and achieved competitive performance. Code: https://github.com/NVlabs/ViR

Figures (7)

• Figure 1: The proposed ViR enables dual parallel and recurrent formulations by using a retention mask. The ViR can be trained in the parallel mode and achieve a competetive performance. The inference can leverage recurrent or chunkwise formulations to improve image throughput and memory efficiency. The effective receptive field and corresponding retention masks in ViR are visualized for a set of patches.
• Figure 2: Overview of the architecture of ViR model. Similar to ViT, Flattened patches are linearly projected into a patch embedding. The position embedding are then added to the patch embedding and a class token is appended to this sequence. The retention encoder comprises of alternating Multi-Head Retention and MLP blocks. The MHR blocks use a causal decay mask. Please see the supplementary materials for detailed information regarding the architecture of H-ViR model.
• Figure 3: Effect of increasing the image resolution on Top-1 accuracy for 1D ViR-B/16 and 2D ViR-B/16 networks.
• Figure 4: Effective receptive field and corresponding masks for: (a)1D retention (b) 2D retention. Cell opacity is based on the decay strength given the distance between the highlighted cell and each of the colored in cells. The 2D formulation achieves shift equivariance, enabling an identical decay factor between successive horizontal and vertical positions. Hence 2D retention is more suitable for finetuning on higher resolutions.
• Figure 5: Effect of image resolution on throughput for ViR-B/16 and ViT-B/16 models. Throughput is measured on an A100 80GB NVIDIA GPU with batch sizes of 16 and 128. For a batch size of 128, the memory is insufficient to process images for both ViT and parallel mode of ViR networks. For $1024 \times 1024$, ViR-B/16 with chunkwise mode is the only configuration that can process images with batch size of 128.