QMamba: Post-Training Quantization for Vision State Space Models

TL;DR

QMamba tackles efficient post-training quantization for vision State Space Models by analyzing activation distributions and designing two targeted quantizers: Long-tailed Skewness Quantization for discrete parameters $\overline{A}_t$ and Temporal Group Quantization for dynamic hidden states $h_t$. It demonstrates that standard uniform quantization underperforms on SSMs and shows substantial accuracy gains over existing PTQ methods, including a 21.0% Top-1 improvement at 4-bit activations on ImageNet. The approach enables hardware-friendly inference via bit-shift operations in the quantized Mamba block and a blockwise fine-tuning pipeline with limited calibration data. These results suggest strong potential for deploying SSM-based vision models on resource-constrained edge devices.

Abstract

State Space Models (SSMs), as key components of Mamaba, have gained increasing attention for vision models recently, thanks to their efficient long sequence modeling capability. Given the computational cost of deploying SSMs on resource-limited edge devices, Post-Training Quantization (PTQ) is a technique with the potential for efficient deployment of SSMs. In this work, we propose QMamba, one of the first PTQ frameworks to our knowledge, designed for vision SSMs based on the analysis of the activation distributions in SSMs. We reveal that the distribution of discrete parameters exhibits long-tailed skewness and the distribution of the hidden state sequence exhibits highly dynamic variations. Correspondingly, we design Long-tailed Skewness Quantization (LtSQ) to quantize discrete parameters and Temporal Group Quantization (TGQ) to quantize hidden states, which reduces the quantization errors. Extensive experiments demonstrate that QMamba outperforms advanced PTQ methods on vision models across multiple model sizes and architectures. Notably, QMamba surpasses existing methods by 21.0% on ImageNet classification with 4-bit activations.

Figures (6)

• Figure 1: The distributions of discrete parameters $\overline{A}_t$ and hidden states $h_t$, which are a part of the state equation ($h_t=\overline{A}_th_{t-1}+\overline{B}_tx_t$) for the input $x_t$ in the SSM operator. The horizontal axis represents the value range. (a) Long-tailed skewed distribution of discrete parameters $\overline{A}_t$. (b) Highly dynamic variation of hidden states $h_t$.
• Figure 2: The statistics of discrete parameters $\overline{A}_t$ and hidden states $h_t$ are observed in the SSM of Vim vim. We visualize the distribution range variations of $\overline{A}_t$ and $h_t$ at different time steps in (a) and (b), where the horizontal axis represents the time dimension. For clarity, we visualize $\overline{A}_t$ and $h_t$ at every fifth time step in the sequence, and the outliers in the boxplot of hidden states $h_t$ are omitted for better visualization. In (c) and (d), we visualize distribution of $\overline{A}_t$ and $h_t$ at two different time step (i.e., $t=30$ and $t=120$), where the horizontal axis represents the value range. The tensor-wise uniform quantization on $\overline{A}_t$ and $h_t$ at different time steps with a single scaling factor results in uniform and same quantization intervals at different time steps.
• Figure 3: Analysis of quantization sensitivity in different SSM activations. The results report the Top-1 and Top-5 accuracy of Vim-S vim on ImageNet. FP denotes the results of the floating-point model. The numbers shown above the bars represent the drop in accuracy compared to FP. Each activation is individually quantized to 4 bits.
• Figure 4: Overview of our QMamba framework. This figure illustrates our quantization framework from the quantized Mamba block to our proposed Long-tailed Skewness Quantization (LtSQ) and Temporal Group Quantization (TGQ) in the quantized SSM. Our LtSQ and TGQ are used for the quantization on discrete parameters $\overline{A}_t$ and hidden states $h_t$, respectively. For a quantized Mamba block, we perform tensor-wise uniform quantization for weights (i.e., $A$ and $D$) and other activations (i.e., $\Delta_t$, $B_t$, $C_t$, and $x_t$) in the SSM operator and for the linear projection and convolution layers in the Mamba block.
• Figure 5: The visualization of our LtSQ and TGQ on discrete parameters $\overline{A}_t$ and hidden states $h_t$ at different time steps in Vim-B (W6A4). (a) The visualization of 4-bit quantization for $\overline{A}_t$ using our LtSQ. (b) The variation of $h_t$ and TGQ scaling factors with time steps. The horizontal axis represents the time step. We show the $h_t$ at every fifth time step and omit outliers in the boxplot for better visualization. (c) The distribution of $h_t$ at different time steps and the corresponding 4-bit quantization intervals using our TGQ.