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

QMamba: Post-Training Quantization for Vision State Space Models

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

https://www.arxiv.org/pdf/2501.13624

摘要

状态空间模型（State Space Models, SSMs）作为Mamba架构的核心组件，近年来因其高效的长序列建模能力而在视觉模型领域受到越来越多的关注。鉴于在资源受限的边缘设备上部署SSMs的计算成本较高，训练后量化（Post-Training Quantization, PTQ）作为一种技术，具有推动SSMs高效部署的潜力。在本研究中，我们提出了QMamba，据我们所知，这是首个专为视觉SSMs设计的PTQ框架之一，其设计基于对SSMs中激活分布的深入分析。我们发现，离散参数的分布呈现出长尾偏态，而隐藏状态序列的分布在时序上表现出高度动态的变化。针对这些特性，我们分别设计了长尾偏态量化方法 （Long-tailed Skewness Quantization, LtSQ）用于量化离散参数，以及时序分组量化方法（Temporal Group Quantization, TGQ）用于量化隐藏状态，从而有效降低了量化误差。

大量实验表明，QMamba在多种模型规模和结构的视觉模型上均优于现有的先进PTQ方法。特别值得一提的是，在ImageNet分类任务中，使用4比特激活的情况下，QMamba相比现有方法性能提升了21.0%。

1. 引言

Mamba [6] 是一种基于状态空间模型（State Space Models, SSMs）[9, 22, 27] 的新型且强大的主干网络，由于其高效的长序列建模能力，已成为研究热点之一。在视觉模型中，基于SSM的模型 [10, 18, 26, 31, 34] 凭借出色的性能和线性时间复杂度取得了令人瞩目的进展，已成为视觉Transformer（Vision Transformers, ViTs）[3, 25] 的有力替代方案。

尽管SSMs在长序列建模方面具有优势，但由于边缘设备在内存和功耗方面的限制，将其部署到各种硬件平台上仍然面临挑战。训练后量化（Post-Training Quantization, PTQ）是解决这一问题的有效方法，它可以在仅使用少量未标注的校准数据集的情况下，将模型的权重和激活值量化为整数，从而减轻内存和功耗负担，便于模型在资源受限的边缘设备上的部署。

然而，现有的PTQ方法主要关注复杂的优化策略 [15, 24, 29]，或针对特定算子（例如ViTs中的Softmax）设计定制化的量化器 [16, 17, 33]。这种研究重点导致对SSMs内部算子的量化分析被忽视，从而在专门针对SSMs的量化方法方面形成了空白。

由于SSMs中的算子与卷积神经网络（CNNs）和视觉Transformer（ViTs）中的算子不同，其中包括状态方程[6]，我们通过对SSMs中激活值分布的分析，揭示了其在量化方面的敏感性和异常值问题。

我们注意到两个对SSMs量化构成挑战的独特特性：

1）如图1a所示，SSMs中离散参数At的分布呈现出长尾偏态分布，即在最大值附近有密集集中，在远离最大值的位置则分布稀疏。这一特性使得均匀量化变得复杂；

2）如图1b所示，SSMs中隐藏状态ht在不同时刻的激活范围具有高度动态性。这种变化性使得在整个序列中使用单一的量化参数变得困难。

在本工作中，基于上述观察结果，我们提出了QMamba，这是首个专为视觉SSMs设计的PTQ框架之一。首先，我们提出了一种长尾偏态量化方法 （Long-tailed Skewness Quantization, LtSQ），以应对离散参数中存在的长尾和偏态分布问题。LtSQ对密集分布的激活值进行非均匀量化，确保量化后的离散参数与量化后的隐藏状态之间的乘法运算可以通过硬件友好的位移操作高效实现。

其次，我们提出了一种时序分组量化方法 （Temporal Group Quantization, TGQ），用于处理隐藏状态序列在不同时刻的高度动态范围。TGQ在时间维度上对隐藏状态序列进行分组并分别量化，从而实现细粒度的、适应隐藏状态变化动态的量化策略。

我们在现有的代表性基于SSM的视觉模型（即Vim [34] 和VMamba [18]）上评估了QMamba的有效性，并在图像分类任务中进行了验证。通过在多种基于SSM的视觉模型和不同比特位宽配置下的全面实验，我们证明了针对SSM算子优化的QMamba在精度方面显著优于当前主流的PTQ方法。特别值得注意的是，在使用4比特激活的情况下，我们的方法在ImageNet分类任务中的Top-1准确率比现有方法提升了21.0%。

本工作的主要贡献如下：

1）据我们所知，QMamba是首个专为视觉模型中SSM算子设计的PTQ框架之一，填补了SSM算子量化方法领域的空白。

2）我们通过揭示基于SSM的模型中激活值的两种分布特性，分析了SSM量化的难点。

3）我们设计了针对长尾偏态分布离散参数的LtSQ方法，以及针对高度动态隐藏状态的TGQ方法，有效克服了SSM量化过程中的挑战。

4）大量实验表明，我们的QMamba在具有代表性的基于SSM的视觉模型上显著优于现有的PTQ方法。

2. 相关工作

2.1. 视觉状态空间模型

状态空间模型（State Space Models, SSMs）因其在建模长序列方面具有线性时间复杂度的潜力而受到越来越多的关注。早期基于SSMs的研究 [7–9, 11] 主要用于处理序列数据，重点在于捕捉长距离依赖关系。在此基础上，Gu 和 Dao [6] 提出了一种新颖的选择性SSM——Mamba，该方法通过将时变参数引入SSM算子中，加入了选择机制，使SSM能够在序列数据的不同时间步有选择地传播或遗忘信息。

在计算机视觉任务中，近期的研究 [10, 18, 26, 34] 将SSMs扩展到处理图像块（image patches），将其视为序列数据以有效建模空间依赖关系。例如，Vim [34] 采用双向Mamba模块构建视觉主干网络来建模视觉表示；VMamba [18] 则通过2D选择扫描模块从多个角度聚合视觉上下文信息。尽管SSMs在长序列建模方面表现出高效性，但在资源受限的边缘设备上部署基于SSM的视觉模型仍是一个有待探索的问题。本研究旨在为SSMs在视觉任务中的量化提供一种有效的方法，推动其在边缘设备上的部署。

2.2. 训练后量化

量化是一种有效的模型压缩技术，它将权重和激活值从浮点数转换为低比特整数，从而减少内存存储和计算消耗。量化方法大致可分为两类：量化感知训练 （Quantization-Aware Training, QAT）[1, 4, 5] 和训练后量化 （Post-Training Quantization, PTQ）[15, 29]。

QAT 在带标签的数据集上联合优化量化参数和模型权重，虽然可以获得较高的精度，但需要大量的训练成本。相比之下，PTQ 方法则专注于仅使用少量未标注的校准数据集来优化量化参数，提供了一种轻量级的替代方案，避免了重新训练的需求。

OMSE [2] 通过最小化权重和激活的量化误差来实现低比特精度推理。Adaround [24] 引入了一种新的权重舍入方法，在低比特位宽下优于传统的最近邻舍入方法。BRECQ [15] 通过顺序重建基本模型模块来提升PTQ性能，在4比特宽度下能够达到与QAT方法相当的表现。Qdrop [29] 则通过在PTQ过程中随机跳过激活值的量化操作进一步提升了PTQ性能，增强了量化模型的鲁棒性。

然而，这些先进的量化方法并非专门为SSMs设计。SSMs中独特的算子结构和动态特性给现有的量化方法带来了挑战。为了填补这一空白，我们提出了首个专为视觉SSMs设计的PTQ方法之一，通过针对SSM算子的定制化设计，确保实现高效的低比特量化。

3. 方法

3.1. 预备知识

选择性状态空间模型（Selective SSMs）的公式。

3.2 SSMs 的量化分析

3.3 长尾偏斜量化

3.4 时间分组量化

3.5 量化Mamba模块

4. 实验与结果 4.1 实验设置

评估设置：在对我们的 QMamba 进行评估时，我们在 ImageNet 分类任务上进行了实验，涉及不同的位宽配置，包括 W8A8（权重和激活均为 8 位量化）、W6A6（6 位量化）以及 W6A4（权重 6 位量化，激活 4 位量化）。我们在基于 SSM 的代表性视觉模型 Vim [34] 和 VMamba [18] 上进行了量化，包括它们各自的 tiny、small 和 base 版本，分别用‘-T’、‘-S’和‘-B’表示。专注于 SSM 中的激活量化，我们为 Mamba 块中线性投影的某些激活设置了特定的位宽。更多关于位宽设置的细节可以在补充材料中找到。

基线方法：我们选择了基于统计的方法（即 MinMax [13]、Percentile [30] 和 OMSE [2]）和基于学习的方法（即 Adaround [24]、BRECQ [15] 和 QDrop [29]）作为比较的基线方法。为了公平比较，所有基于学习的方法的学习设置与我们的 QMamba 保持一致。

4.2 比较结果

我们在 ImageNet 分类任务上对 Vim 和 VMamba 的多个量化版本进行了 QMamba 的评估。如表 1 所示，我们的 QMamba 在 Vim 和 VMamba 上始终实现了最高的 Top-1 和 Top-5 准确率。我们注意到，激活值的位宽越低，我们的 QMamba 相对于基线方法的优势就越明显。例如，在 Vim-T（W6A6）上，我们的 QMamba 的 Top-1 和 Top-5 准确率比 QDrop 高出 5.5% / 6.8%（57.9% / 82.0% 对比 52.4% / 75.2%），而在 VMamba-T（W6A4）上，我们的 QMamba 的 Top-1 和 Top-5 准确率比 QDrop 高出 3.0% / 2.0%（54.9% / 79.2% 对比 51.9% / 77.2%）。值得注意的是，我们的 QMamba 在 Vim-S（W6A4）上大幅超越了所有基线方法，例如，Top-1 和 Top-5 准确率比 QDrop 高出 21.0% / 26.4%（45.9% / 69.7% 对比 24.9% / 43.3%）。这是因为对所有激活值应用张量级均匀量化会导致较大的量化误差。我们提出的 LtSQ 有效地避免了这一问题。

5. 消融分析
5.1 缩放因子的初始化

如第 3.2 节所述，SSM 中的激活值对量化非常敏感，这表明缩放因子的初始化对我们提出的 QMamba 来说非常重要。我们在表 2 中报告了在不微调任何缩放因子的情况下，使用我们 QMamba 初始化模型的结果。我们分别使用 MinMax [13]、Percentile [30] 和 OMSE [2] 方法对缩放因子进行初始化。其中 Percentile 方法使用校准数据集上观察到的第 1 百分位和第 99 百分位值来初始化缩放因子。

在此，我们得出以下三个结论：

1）MinMax 不适合用于 SSM 中激活值的初始化方法 。如表 2 所示，所有使用 MinMax 初始化缩放因子的模型表现都显著差于使用 Percentile 和 OMSE 初始化的模型。例如，在 Vim-S（W6A6）上，使用 MinMax 初始化会导致性能下降 57.4%（相比于使用 Percentile 初始化）。

2）Percentile 是一种针对 SSM 中激活值的稳健初始化方法 。例如，使用 Percentile 初始化时，Vim-S（W6A6）的表现比使用 OMSE 高出 23.0%，而 VMamba-S（W6A4）的表现则高出 16.7%。

3）OMSE 不适合用于 SSM 中低比特激活量化的初始化方法 。当 SSM 的激活值被量化为高比特位宽时，使用 OMSE 的初始化效果与使用 Percentile 相当。然而，在低比特位宽的情况下，OMSE 的初始化效果通常不如 Percentile。其原因是 OMSE 的初始化过程是基于最小化 MSE 损失函数的，而在低比特位宽情况下，SSM 激活值中存在大规模异常值，这会影响 OMSE 的初始化效果。

5.2 LtSQ 的有效性

表 3 中列出的 Top-1 准确率结果展示了我们提出的 LtSQ 的有效性。例如，仅使用我们的 LtSQ 对 Vim-S、Vim-B 和 VMamba-B 进行量化时，其 Top-1 准确率分别比相应的基线模型高出 13.4%、1.5% 和 1.9%。

当我们联合使用 LtSQ 和我们的 TGQ 时，量化模型的性能提升更加明显。特别地，使用 LtSQ 与 TGQ 联合量化 Vim-S 时，其 Top-1 准确率比基线高出 21.0%（45.9% vs. 24.9%）。即使仅使用 LtSQ，Vim-S 的准确率也比基线高出 13.4%（38.3% vs. 24.9%）。我们将这一提升归因于在使用 LtSQ 相比使用均匀量化器时，量化后的 Vim-S 在初始化阶段表现更好，从而有助于后续缩放因子的微调过程。

如图 5a 所示，我们的 LtSQ 在长尾分布密集区域使用了更细粒度的量化区间，从而避免了当离散参数 At 被量化为低位宽时出现的不良初始化问题。此外，我们还设置了 λ = 0.9，并尝试了不同的偏态边界 α 进行实验。如表 4 所示，当 α = 0.9 时，Top-1 准确率达到最佳结果。

5.3 TGQ 的有效性

我们在表 3 中列出了相关结果，以评估我们提出的 TGQ 的有效性。例如，在 Vim-B 上，仅使用 TGQ 时其 Top-1 准确率比基线提升了 1.6%，而将 TGQ 与 LtSQ 联合使用时，提升幅度进一步提高到 2.9%。这说明无论是单独使用还是与 LtSQ 结合使用，我们的 TGQ 都能够有效提升量化模型的性能。

为进一步分析 TGQ 的有效性，我们在图 5b 中可视化了 Vim 上 TGQ 的缩放因子。TGQ 按时间步顺序将隐藏状态 ht 分成若干组，并使用相应的缩放因子进行量化。图中每个红色点代表一组对应的缩放因子。在这里，我们将 TGQ 的组长度 λ 设为 10。如图 5b 所示，缩放因子的变化趋势与隐藏状态分布范围的变化趋势是一致的。

此外，如图 5c 所示，我们使用 TGQ 可视化了在两个不同时间步上隐藏状态 ht 的量化区间。其中，分布范围较大的隐藏状态 h120 使用较大的量化区间进行量化，而分布范围较小的隐藏状态 h30 则使用较小的量化区间进行量化。需要注意的是，根据公式（3），量化区间与对应的缩放因子成正比。这表明我们的 TGQ 能够灵活地对不同时刻动态变化的隐藏状态 ht 进行量化。

另外，我们还设置了 α = 0.9，并尝试了不同的组长度 λ 进行实验。如表 4（行编号：3、5 和 6）所示，对于 Vim-B 和 VMamba-B，当 λ 设置为 10 时，TGQ 能够取得最佳效果。

5.4 理论效率结果

为了评估我们提出的 QMamba 在理论效率方面的表现，我们按照 [18, 19, 21] 的方法计算了量化后的 Vim 的计算 FLOPs（浮点运算次数）和存储大小。如图 6 所示，相比浮点模型，使用 QMamba 量化的模型能够在性能、存储大小和计算 FLOPs 之间实现更好的平衡。

例如，当将 Vim-B 量化为 6 位时，QMamba 最多可节省 80% 的存储空间和 75% 的 FLOPs，同时仍保持较高的 Top-1 准确率 75.8%（原始浮点模型为 81.9%）。

6 结论
在本研究中，我们首次提出了 QMamba，这是首个专为基于状态空间模型（SSM）的视觉模型设计的后训练量化（PTQ）框架之一。我们的 QMamba 解决了由于 SSM 中独特算子特性所带来的量化挑战。通过对 SSM 中离散参数和隐藏状态的分布进行分析，我们发现了量化过程中两个关键挑战：一是长尾偏态分布的离散参数，二是随时间步动态变化的隐藏状态。针对这些问题，我们分别引入了 LtSQ 来处理长尾偏态分布，以及 TGQ 来处理不同时刻隐藏状态的动态范围。

大量实验结果表明，我们的 QMamba 在具有代表性的基于 SSM 的视觉模型上取得了优异的量化效果，使得这些模型能够在资源受限的边缘设备上高效部署。

https://www.arxiv.org/pdf/2501.13624