QMamba: Post-Training Quantization for Vision State Space Models

📄 arXiv: 2501.13624v1 📥 PDF

作者: Yinglong Li, Xiaoyu Liu, Jiacheng Li, Ruikang Xu, Yinda Chen, Zhiwei Xiong

分类: cs.CV

发布日期: 2025-01-23

💡 一句话要点

QMamba：面向视觉状态空间模型的后训练量化框架

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 后训练量化 状态空间模型 视觉模型 模型压缩 边缘计算

📋 核心要点

  1. 现有视觉状态空间模型部署成本高，后训练量化是降低成本的有效手段，但现有PTQ方法效果不佳。
  2. QMamba通过分析SSM激活分布，针对离散参数和隐藏状态序列的特性，分别设计了长尾偏斜量化和时间分组量化。
  3. 实验表明，QMamba在多种视觉模型和数据集上显著优于现有PTQ方法，例如在ImageNet分类上提升高达21.0%。

📝 摘要（中文）

状态空间模型(SSM)作为Mamba的关键组成部分，由于其高效的长序列建模能力，近年来在视觉模型中受到越来越多的关注。考虑到在资源受限的边缘设备上部署SSM的计算成本，后训练量化(PTQ)是一种具有高效部署SSM潜力的技术。本文提出了QMamba，据我们所知，它是首批基于SSM中激活分布分析而设计的视觉SSM的PTQ框架之一。我们发现离散参数的分布表现出长尾偏斜，而隐藏状态序列的分布表现出高度动态的变化。相应地，我们设计了长尾偏斜量化(LtSQ)来量化离散参数，以及时间分组量化(TGQ)来量化隐藏状态，从而减少量化误差。大量的实验表明，QMamba在多种模型大小和架构的视觉模型上优于先进的PTQ方法。值得注意的是，QMamba在具有4位激活的ImageNet分类上超过了现有方法21.0%。

🔬 方法详解

问题定义：论文旨在解决视觉状态空间模型（Vision SSMs）在资源受限的边缘设备上部署时计算成本过高的问题。现有的后训练量化（PTQ）方法在应用于Vision SSMs时，由于其特殊的激活分布特性，量化效果不佳，导致精度损失较大。因此，需要一种专门为Vision SSMs设计的PTQ框架，以实现高效且准确的模型量化。

核心思路：QMamba的核心思路是针对Vision SSMs中离散参数和隐藏状态序列的特殊激活分布特性，分别设计定制化的量化策略。具体来说，论文观察到离散参数的分布呈现长尾偏斜，而隐藏状态序列的分布具有高度动态变化。因此，需要设计能够适应这些特性的量化方法，以最小化量化误差。

技术框架：QMamba的整体框架包括两个主要模块：长尾偏斜量化（LtSQ）和时间分组量化（TGQ）。首先，对SSM中的离散参数进行分析，利用LtSQ来处理其长尾偏斜分布。然后，对隐藏状态序列进行分析，利用TGQ来处理其动态变化。最后，将量化后的模型部署到边缘设备上。

关键创新：QMamba的关键创新在于针对Vision SSMs的特殊激活分布，提出了两种定制化的量化策略：LtSQ和TGQ。LtSQ通过考虑长尾偏斜分布，更准确地量化离散参数。TGQ通过对隐藏状态序列进行时间分组，并对每个组进行独立的量化，从而更好地适应其动态变化。与现有PTQ方法相比，QMamba能够显著降低量化误差，提高量化模型的精度。

关键设计：LtSQ的关键设计在于如何有效地处理长尾偏斜分布。具体实现细节未知，但推测可能采用了非对称量化或者针对长尾部分进行特殊处理的策略。TGQ的关键设计在于如何确定时间分组的大小和量化参数。具体实现细节未知，但推测可能采用了动态规划或者启发式搜索的方法来优化分组策略，并根据每个组的统计特性来确定量化参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

QMamba在ImageNet分类任务上取得了显著的性能提升。使用4位激活量化时，QMamba超过了现有PTQ方法21.0%。实验结果表明，QMamba能够有效地降低Vision SSMs的量化误差，并保持较高的模型精度。此外，QMamba在多种模型大小和架构的视觉模型上都表现出优越的性能，证明了其通用性和有效性。

🎯 应用场景

QMamba的潜在应用领域包括智能手机、无人机、机器人等资源受限的边缘设备上的图像分类、目标检测、图像分割等视觉任务。通过降低Vision SSMs的计算成本和模型大小，QMamba可以实现更高效、更实时的视觉应用，并促进边缘智能的发展。未来，QMamba可以进一步扩展到其他类型的SSM和视觉模型，并与其他模型压缩技术相结合，以实现更高的压缩率和更好的性能。

📄 摘要（原文）

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.