Vision Retention Networks (ViR)

• Vision Retention Networks (ViR) are neural architectures that integrate retention operators into vision models to efficiently capture local context and reduce computational costs compared to traditional self-attention.
• They employ dual parallel and recurrent retention operators with convolutional-like attention masking, enhancing performance on image classification and dense prediction tasks.
• ViR scales through recurrent, chunkwise, and 2D shift-equivariant formulations, achieving state-of-the-art results on benchmarks like ImageNet and COCO while supporting efficient inference.

Vision Retention Networks (ViR) are a class of neural network architectures that incorporate the retention mechanism—originally developed for efficient sequence modeling in language and adopted from RetNet—into computer vision backbones for image classification and dense prediction tasks. ViR architectures address the computational inefficiency of self-attention in Vision Transformers (ViTs), while enabling explicit local inductive bias and efficient inference. They achieve this by introducing dual parallel and recurrent retention operators, convolutional-like attention masking, and novel parameter-efficient locality modeling strategies (Li et al., 2023, Hatamizadeh et al., 2023).

1. Foundational Concepts and Motivation

ViT architectures employ self-attention to capture long-range dependencies among image patches, but the quadratic computational cost of $O(N^2)$ in sequence length $N$ imposes limitations on memory and throughput, especially for high-resolution images and generative (autoregressive) applications. In natural language processing, RetNet replaces softmax attention with a retention operator leveraging exponential-decay distance priors, achieving efficient parallel training and swift recurrent inference.

In computer vision, inductive priors such as locality and weight sharing, exploited by CNNs, are lacking in standard ViTs, causing their underperformance on small/mid-sized datasets and resource-constrained regimes. RetNet-inspired architectures introduce explicit local bias into the visual domain by modulating the attention matrix via learnable or convolutional masks, facilitating the modeling of local context in a manner analogous to spatial convolution (Li et al., 2023).

2. Retention Mechanism: Parallel, Recurrent, and Chunkwise Formulations

The retention operator replaces multi-head self-attention in ViTs. Token embeddings $\mathbf{Z} \in \mathbb{R}^{N \times D}$ are mapped via learned projections to query, key, and value matrices. The core retention can be expressed:

• Recurrent Formulation: Each token sequentially updates a memory state $\mathbf{s}_n$ with exponential decay $\gamma \in (0,1)$:

$$\mathbf{s}_n = \gamma\,\mathbf{s}_{n-1} + \mathbf{k}_n^\top\,\mathbf{v}_n, \quad \mathrm{Ret}(\mathbf{Z}_n) = \mathbf{q}_n\,\mathbf{s}_n$$

This admits $O(N)$ time and $O(1)$ additional memory for inference.

• Parallel (Masked) Formulation: All tokens are processed concurrently with a lower-triangular decay mask $\mathbf{M} \in \mathbb{R}^{N \times N}$:

$$\mathrm{Ret}(\mathbf{Z}) = \left(\frac{\mathbf{q}\,\mathbf{k}^\top}{\sqrt{D_h}} \odot \mathbf{M}\right)\,\mathbf{v}, \quad M_{ij} = \begin{cases} \gamma^{i-j}, & i \geq j\ 0, & i < j \end{cases}$$

This matches the computational structure of attention during training and enables batch-parallelism.

• Chunkwise (Hybrid) Formulation: The sequence is partitioned into chunks of size $C$. Parallel retention is applied within each chunk, while a summary state is carried between chunks, reducing complexity to $O(NC)$ and memory to $O(C^2)$.
• 2D Shift-Equivariant Retention: To better align with image structure and avoid scanline bias, retention is extended to a 2D grid, ensuring equivariance to spatial shifts. The decay applies identically over horizontal and vertical axes, resulting in improved accuracy at high resolutions (Hatamizadeh et al., 2023).

3. Locality Modeling: ELM and GMM Attention Masks

RetNet-derived ViR variants incorporate explicit local inductive bias through attention masking:

• Element-wise Learnable Mask (ELM): For each attention head, an $N \times N$ mask $M_{ELM}$ is learned, modulating the attention matrix by Hadamard product:

$A' = A \odot M_{ELM}$

While fully expressive, ELM introduces $O(N^2)$ parameters per head, leading to increased memory, compute requirements, and risk of overfitting on smaller datasets.

• Gaussian Mixture Mask (GMM): Locality is captured more efficiently by modeling the attention mask as a mixture of $K$ isotropic Gaussians with learnable amplitudes $\alpha_k$ and widths $\sigma_k$:

$$M_{GMM}[i, j] = \sum_{k=1}^K \alpha_k \exp\left[-\frac{(x_{ij}^2 + y_{ij}^2)}{2\sigma_k^2 + \epsilon}\right]$$

This drastically reduces overhead (only $2K$ new parameters per head per layer; $K=5$ is sufficient), incurs negligible additional compute, and improves performance on local pattern-dominated small- and mid-scale datasets (Li et al., 2023).

Inserting these masks—either ELM or GMM—before/after softmax in attention layers, ViR backbones (ViT, Swin, CaiT, PiT, etc.) benefit from locality without sacrificing global modeling.

4. Architectural Integration and Variants

• ViR-Backbone: Images are split into non-overlapping patches; tokens are embedded and position-encoded. A [class] token is appended for classification tasks.
• Attention Replacement: The attention mechanism is replaced by Multi-Head Retention (MHR), incorporating the decay mask or recurrent carry-state, alternating with standard MLP and LayerNorm blocks.
• GMM Insertion: In ViT variants, the GMM mask is applied to the attention scores before softmax. For Swin-style local windows, the mask is sized per window. For CaiT, the GMM is used in standard self-attention (not the final class-token block).
• Pooling Strategy: The class-token can be dropped and global average-pooling applied on output tokens, maintaining a square attention matrix (Li et al., 2023).

5. Computational Efficiency and Scalability

ViR achieves a favorable throughput-to-accuracy trade-off by supporting both highly parallel training and efficient sequential or chunkwise inference:

• Self-attention: $O(N^2)$ time and memory.
• Recurrent retention: $O(N)$ time, $O(1)$ extra memory.
• Chunkwise retention: $O(NC)$ time, $O(C^2)$ memory, where $C \ll N$.
• 2D shift-equivariant retention: Offers additional gains for large images due to spatially uniform decay.

Chunkwise mode enables large-batch, high-resolution processing—e.g., ViR-B/16, chunkwise, at $1024 \times 1024$ with batch size 128, outscaling baseline ViTs in throughput and memory—while inference in recurrent mode is amenable to streaming or autoregressive tasks and is more than $3\times$ faster than standard ViTs for equivalent settings (Hatamizadeh et al., 2023).

6. Empirical Evaluation and Benchmarks

ViR demonstrates strong empirical results across a range of datasets, model sizes, and tasks:

• Small Data Regimes: On CIFAR-10, CIFAR-100, SVHN, and Tiny-ImageNet, GMM-masked ViT backbones show consistent and significant improvements in top-1 accuracy with negligible computational overhead, outperforming both unmodified ViT and CNN hybrids such as Swin, PiT, T2T-ViT, and CaiT.
• ImageNet Scale: For classification, ViR and hybrid HViR models deliver competitive or superior throughput and accuracy relative to ViT and Swin:
  • HViR-2 (56M params) achieves 83.3% top-1, exceeding Swin-S at 83.2%.
  • HViR-3 scales to 84.6% top-1.
  • ViR-L/14 achieves up to 86.1% (448×448 input).
• Detection/Segmentation: HViR-1 achieves 51.7 AP$^{box}$ and 44.1 AP$^{mask}$ on COCO; mIoU 47.0 on ADE20K, outperforming Swin-T at 50.4/43.7 and 44.5 mIoU (Hatamizadeh et al., 2023).
• Ablation: GMM performance increases with number of mixture kernels $K$ up to $K\approx 5$, plateauing thereafter (Li et al., 2023).

Example gains from GMM-masked variants:

Model	CIFAR-10	CIFAR-100	SVHN	Tiny-ImageNet
ViT-9	93.65	75.36	97.93	59.89
GMM-ViT	95.06	77.81	98.01	62.27
Swin	95.26	77.88	97.89	60.45
GMM-Swin	95.39	78.26	97.90	61.03

7. Strengths, Limitations, and Future Research

Strengths:

• Dual parallel/recurrent equivalence enables efficient training and deployment.
• Chunkwise formulation decouples peak memory from sequence length.
• GMM masking provides locality and strong small-data generalization at negligible parameter cost.
• 2D retention offers shift-equivariance and superior scaling as image resolution increases.
• Native support for autoregressive and generative vision use cases.

Limitations:

• Purely recurrent mode remains $O(N)$ sequential steps at inference.
• Efficacy of retention relative to full self-attention on intricate long-range dependencies warrants further study.
• Large-scale generative vision modeling (autoregressive image synthesis, video prediction) remains untested (Hatamizadeh et al., 2023).

Future Directions:

• Learnable relative position encoding integration into retention (e.g., xPos, RoPE).
• Hybridization with CNNs for increased data efficiency.
• Extension to multi-modal architectures.
• Self-supervised learning with retention.
• Broader autoregressive/compositional vision tasks (Hatamizadeh et al., 2023).

• "Toward a Deeper Understanding: RetNet Viewed through Convolution" (Li et al., 2023)
• "ViR: Towards Efficient Vision Retention Backbones" (Hatamizadeh et al., 2023)