SparX: A Sparse Cross-Layer Connection Mechanism for Hierarchical Vision Mamba and Transformer Networks

Abstract

Due to the capability of dynamic state space models (SSMs) in capturing long-range dependencies with linear-time computational complexity, Mamba has shown notable performance in NLP tasks. This has inspired the rapid development of Mamba-based vision models, resulting in promising results in visual recognition tasks. However, such models are not capable of distilling features across layers through feature aggregation, interaction, and selection. Moreover, existing cross-layer feature aggregation methods designed for CNNs or ViTs are not practical in Mamba-based models due to high computational costs. Therefore, this paper aims to introduce an efficient cross-layer feature aggregation mechanism for vision backbone networks. Inspired by the Retinal Ganglion Cells (RGCs) in the human visual system, we propose a new sparse cross-layer connection mechanism termed SparX to effectively improve cross-layer feature interaction and reuse. Specifically, we build two different types of network layers: ganglion layers and normal layers. The former has higher connectivity and complexity, enabling multi-layer feature aggregation and interaction in an input-dependent manner. In contrast, the latter has lower connectivity and complexity. By interleaving these two types of layers, we design a new family of vision backbone networks with sparsely cross-connected layers, achieving an excellent trade-off among model size, computational cost, memory cost, and accuracy in comparison to its counterparts. For instance, with fewer parameters, SparX-Mamba-T improves the top-1 accuracy of VMamba-T from 82.5\% to 83.5\%, while SparX-Swin-T achieves a 1.3\% increase in top-1 accuracy compared to Swin-T. Extensive experimental results demonstrate that our new connection mechanism possesses both superior performance and generalization capabilities on various vision tasks.

Accepted by AAAI 2025. Code is available at https://github.com/LMMMEng/SparX

References in corpus (20)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• PVT v2: Improved Baselines with Pyramid Vision Transformer
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• CoAtNet: Marrying Convolution and Attention for All Data Sizes
• Twins: Revisiting the Design of Spatial Attention in Vision Transformers
• Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
• Conditional Positional Encodings for Vision Transformers
• EfficientFormer: Vision Transformers at MobileNet Speed
• CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows
• More ConvNets in the 2020s: Scaling up Kernels Beyond 51x51 using Sparsity
• CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention
• Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs
• CMT: Convolutional Neural Networks Meet Vision Transformers
• UniRepLKNet: A Universal Perception Large-Kernel ConvNet for Audio, Video, Point Cloud, Time-Series and Image Recognition
• UniFormer: Unifying Convolution and Self-attention for Visual Recognition
• MPViT: Multi-Path Vision Transformer for Dense Prediction
• DaViT: Dual Attention Vision Transformers
• Shunted Self-Attention via Multi-Scale Token Aggregation
• Slide-Transformer: Hierarchical Vision Transformer with Local Self-Attention
• Fcaformer: Forward Cross Attention in Hybrid Vision Transformer