PTQ4VM: Post-Training Quantization for Visual Mamba

📄 arXiv: 2412.20386v2 📥 PDF

作者: Younghyun Cho, Changhun Lee, Seonggon Kim, Eunhyeok Park

分类: cs.CV, cs.LG

发布日期: 2024-12-29 (更新: 2025-04-07)

备注: Accepted at WACV 2025 (oral presentation)

🔗 代码/项目: GITHUB

💡 一句话要点

PTQ4VM：面向Visual Mamba的后训练量化方法，显著提升推理速度。

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

关键词: 后训练量化 Visual Mamba 模型加速 视觉任务 深度学习 边缘计算 模型优化

📋 核心要点

  1. Visual Mamba虽然高效，但对量化敏感，导致性能提升受限，主要挑战在于token方差、通道异常值和激活长尾分布。
  2. PTQ4VM通过Per-Token Static量化和Smoothing Scale与Step Size联合学习，解决Visual Mamba的量化难题。
  3. 实验表明，PTQ4VM在保证精度损失可忽略不计的情况下，实现了高达1.83倍的GPU加速，具有实际应用价值。

📝 摘要（中文）

Visual Mamba是一种将选择性空间状态模型Mamba扩展到视觉任务的方法。它以固定的顺序依次处理图像tokens，累积信息以生成输出。尽管它在各种任务中以低计算成本提供高质量输出而越来越受欢迎，但Visual Mamba极易受到量化的影响，这使得进一步的性能改进具有挑战性。我们的分析表明，Visual Mamba中固定的token访问顺序引入了独特的量化挑战，我们将其归类为三个主要问题：1) token间的方差，2) 通道间的异常值，以及3) 激活的长尾分布。为了应对这些挑战，我们提出了Visual Mamba的后训练量化（PTQ4VM），它引入了两个关键策略：Per-Token Static（PTS）量化和Smoothing Scale与Step Size的联合学习（JLSS）。据我们所知，这是第一个关于Visual Mamba的量化研究。PTQ4VM可以应用于各种Visual Mamba骨干网络，在不到15分钟内将预训练模型转换为量化格式，而不会显着降低质量。在大型分类和回归任务上的大量实验证明了其有效性，与FP16相比，在GPU上实现了高达1.83倍的加速，而精度损失可忽略不计。我们的代码可在https://github.com/YoungHyun197/ptq4vm获得。

🔬 方法详解

问题定义：Visual Mamba在量化过程中面临显著的性能下降，这是由于其固定的token处理顺序导致了token间方差大、通道间存在异常值以及激活值呈现长尾分布等问题。现有的量化方法无法有效解决这些问题，限制了Visual Mamba在资源受限设备上的部署和应用。

核心思路：PTQ4VM的核心思路是通过针对Visual Mamba特性的量化策略，缓解量化带来的精度损失。具体来说，通过Per-Token Static (PTS) 量化来解决token间方差问题，并利用Joint Learning of Smoothing Scale and Step Size (JLSS) 来处理通道异常值和激活长尾分布。这种针对性的设计旨在最小化量化误差，从而保持模型性能。

技术框架：PTQ4VM的整体框架包括两个主要阶段：首先，对Visual Mamba模型进行离线分析，识别出需要特殊处理的层和激活。然后，应用PTS量化和JLSS策略对模型进行量化。PTS量化对每个token采用独立的量化参数，以适应token间的差异。JLSS策略通过联合学习平滑尺度和步长，优化量化范围，从而更好地处理异常值和长尾分布。

关键创新：PTQ4VM的关键创新在于其针对Visual Mamba量化特性的定制化设计。与传统的量化方法不同，PTQ4VM充分考虑了Visual Mamba中固定的token处理顺序带来的影响，并提出了相应的解决方案。PTS量化和JLSS策略的结合，能够有效地缓解量化带来的精度损失，从而实现更高的量化性能。

关键设计：PTS量化的关键在于为每个token选择合适的量化参数，这可以通过统计每个token的激活值分布来实现。JLSS策略的关键在于设计合适的损失函数，以联合优化平滑尺度和步长。论文中可能使用了诸如KL散度或MSE等损失函数来衡量量化前后的激活值分布差异，并通过梯度下降等优化算法来更新参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PTQ4VM在大型分类和回归任务上表现出色，与FP16相比，在GPU上实现了高达1.83倍的加速，而精度损失可忽略不计。例如，在ImageNet分类任务上，使用PTQ4VM量化后的Visual Mamba模型在保持精度基本不变的情况下，推理速度显著提升。这些结果验证了PTQ4VM的有效性，并表明其具有实际应用价值。

🎯 应用场景

PTQ4VM具有广泛的应用前景，可用于加速各种基于Visual Mamba的视觉任务，如图像分类、目标检测、语义分割等。该方法尤其适用于资源受限的边缘设备，例如移动设备、嵌入式系统和物联网设备，能够以较低的计算成本实现高性能的视觉推理。此外，PTQ4VM还可以促进Visual Mamba在工业自动化、智能安防和自动驾驶等领域的应用。

📄 摘要（原文）

Visual Mamba is an approach that extends the selective space state model, Mamba, to vision tasks. It processes image tokens sequentially in a fixed order, accumulating information to generate outputs. Despite its growing popularity for delivering high-quality outputs at a low computational cost across various tasks, Visual Mamba is highly susceptible to quantization, which makes further performance improvements challenging. Our analysis reveals that the fixed token access order in Visual Mamba introduces unique quantization challenges, which we categorize into three main issues: 1) token-wise variance, 2) channel-wise outliers, and 3) a long tail of activations. To address these challenges, we propose Post-Training Quantization for Visual Mamba (PTQ4VM), which introduces two key strategies: Per-Token Static (PTS) quantization and Joint Learning of Smoothing Scale and Step Size (JLSS). To the our best knowledge, this is the first quantization study on Visual Mamba. PTQ4VM can be applied to various Visual Mamba backbones, converting the pretrained model to a quantized format in under 15 minutes without notable quality degradation. Extensive experiments on large-scale classification and regression tasks demonstrate its effectiveness, achieving up to 1.83x speedup on GPUs with negligible accuracy loss compared to FP16. Our code is available at https://github.com/YoungHyun197/ptq4vm.