Table of Contents
Fetching ...

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

Meng Lou, Yunxiang Fu, Yizhou Yu

TL;DR

SparX introduces a sparse cross-layer connection mechanism for Mamba- and Transformer-based vision backbones, inspired by retinal ganglion cells. By using ganglion and normal layers together and a Dynamic Multi-layer Channel Aggregator (DMCA) within a cross-layer sliding window, SparX enables efficient cross-layer feature interaction and distillation. Empirical results across ImageNet-1K, COCO, and ADE20K show consistent gains in top-1 accuracy and dense prediction metrics, with favorable parameter and throughput trade-offs compared to existing Mamba and ViT backbones. The approach demonstrates strong generalization and versatility, improving both classification and dense prediction tasks while remaining computationally efficient.

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.

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

TL;DR

SparX introduces a sparse cross-layer connection mechanism for Mamba- and Transformer-based vision backbones, inspired by retinal ganglion cells. By using ganglion and normal layers together and a Dynamic Multi-layer Channel Aggregator (DMCA) within a cross-layer sliding window, SparX enables efficient cross-layer feature interaction and distillation. Empirical results across ImageNet-1K, COCO, and ADE20K show consistent gains in top-1 accuracy and dense prediction metrics, with favorable parameter and throughput trade-offs compared to existing Mamba and ViT backbones. The approach demonstrates strong generalization and versatility, improving both classification and dense prediction tasks while remaining computationally efficient.

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.
Paper Structure (22 sections, 1 equation, 5 figures, 17 tables)

This paper contains 22 sections, 1 equation, 5 figures, 17 tables.

Figures (5)

  • Figure 1: Performance comparison between SparX-Mamba and other methods on various vision tasks.
  • Figure 2: A concrete example of proposed SparX with S=2 and M=2.
  • Figure 3: An illustration of the proposed DMCA. The red font indicates the dimensions in matrix multiplications.
  • Figure 4: Visualization of the effective receptive fields of representative models.
  • Figure 5: CKA analysis of model variants