V"Mean"ba: Visual State Space Models only need 1 hidden dimension

TL;DR

This work addresses the GPU underutilization in VMamba-based Vision State Space Models caused by the selective scan bottleneck. It introduces VMeanba, a training-free activation compression that reduces input channels from $D$ to $1$ by applying a mean transform $T$ along the inner channel axis, with an inverse transform $T^{-1}$ to recover outputs, leveraging the observed low cross-channel variance in activations. A layer-impact score guides a layer-selection pipeline to replace $K$ Mamba blocks with VMeanba blocks, achieving substantial FLOP reductions without large accuracy losses, and complementary benefits when combined with pruning. Empirical results on ImageNet-1k and ADE20k show end-to-end speedups up to $1.12\times$ and kernel-level gains up to $\approx 293\times$, while maintaining accuracy within $3\%$, highlighting a practical path to efficient SSM-based vision models; future directions include kernel fusion and extending to additional CV tasks.

Abstract

Vision transformers dominate image processing tasks due to their superior performance. However, the quadratic complexity of self-attention limits the scalability of these systems and their deployment on resource-constrained devices. State Space Models (SSMs) have emerged as a solution by introducing a linear recurrence mechanism, which reduces the complexity of sequence modeling from quadratic to linear. Recently, SSMs have been extended to high-resolution vision tasks. Nonetheless, the linear recurrence mechanism struggles to fully utilize matrix multiplication units on modern hardware, resulting in a computational bottleneck. We address this issue by introducing \textit{VMeanba}, a training-free compression method that eliminates the channel dimension in SSMs using mean operations. Our key observation is that the output activations of SSM blocks exhibit low variances across channels. Our \textit{VMeanba} leverages this property to optimize computation by averaging activation maps across the channel to reduce the computational overhead without compromising accuracy. Evaluations on image classification and semantic segmentation tasks demonstrate that \textit{VMeanba} achieves up to a 1.12x speedup with less than a 3\% accuracy loss. When combined with 40\% unstructured pruning, the accuracy drop remains under 3\%.

Figures (4)

• Figure 1: The GPU kernel time of each operation in a VMamba block. The latency is measured using feature maps with an input resolution of $224 \times 224$. We rank the kernels by their latency and highlights the top-5 time-consuming kernels on the bar chart. The selective scan operation is one of the major contributors in the VMamba block.
• Figure 2: Overview of the VMeanba block. VMeanba reduced the channel dimension of the inputs to the associated scan operation by applying a transform $T$, thereby simplifying the computation. The proposed VMeanba components are highlighted in red, while the original selective scan components are shown in blue and green, with the green block indicating the main area of optimization.
• Figure 3: The figure illustrates the distribution of inner dimension values of $y_{layer}$ across various data points as a function of sequence length. Notably, the distribution remains remarkably consistent across different data points for identical $l$ values, as indicated by the arrows. The distribution for $l=195$ , shown on the right, provides further evidence of this concentration.
• Figure 4: Accuracy versus K Analysis on classification and segmentation tasks by using VMeanba. This figure illustrates the trade-off between the value of $K$ and the associated accuracy drop. By carefully selecting an appropriate $K$, the model's accuracy can be largely preserved.