QuarterMap: Efficient Post-Training Token Pruning for Visual State Space Models

📄 arXiv: 2507.09514v1 📥 PDF

作者: Tien-Yu Chi, Hung-Yueh Chiang, Diana Marculescu, Kai-Chiang Wu

分类: cs.CV, cs.AI

发布日期: 2025-07-13

备注: Accepted by Efficient Systems for Foundation Models Workshop at the International Conference on Machine Learning (ICML) 2025

💡 一句话要点

QuarterMap：为视觉状态空间模型设计的高效后训练Token剪枝方法

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

关键词: 状态空间模型 Token剪枝 后训练量化 模型加速 视觉模型

📋 核心要点

  1. VMamba等视觉状态空间模型在四向扫描中存在空间冗余，限制了其效率。
  2. QuarterMap通过在扫描前剪枝冗余激活，并使用最近邻上采样恢复维度，实现加速。
  3. 实验表明，QuarterMap在VMamba和MedMamba上均能显著提高吞吐量，同时保持精度。

📝 摘要（中文）

状态空间模型（SSM）通过利用线性递归来降低Transformer的二次复杂度。最近，VMamba作为一种强大的基于SSM的视觉骨干网络而出现，但其四向扫描中的空间冗余仍然是瓶颈。我们提出了QuarterMap，一种后训练激活剪枝方法，它在扫描之前移除冗余的空间激活，并通过最近邻上采样恢复维度。我们的方法在不重新训练的情况下提高了吞吐量。在ImageNet-1K上，QuarterMap在VMamba上实现了高达11%的加速，而精度下降不到0.9%，并在ADE20K分割上产生了类似的收益。除了VMamba，我们还在MedMamba（一种共享相同四向扫描结构的领域特定模型）上验证了QuarterMap，它在多个医学成像任务中始终提高吞吐量，同时保持精度。与ToMe等Token合并方法相比，QuarterMap专为SSM设计，避免了代价高昂的合并-解合并操作。我们的方法提供了一种即插即用的工具，用于部署时提高效率，而不会影响可迁移性。

🔬 方法详解

问题定义：VMamba等基于状态空间模型的视觉骨干网络，虽然在计算复杂度上优于Transformer，但其四向扫描机制引入了大量的空间冗余，导致计算效率瓶颈。现有Token合并方法（如ToMe）虽然可以减少Token数量，但需要进行代价高昂的合并-解合并操作，不适用于SSM结构。

核心思路：QuarterMap的核心思想是在不影响模型精度的情况下，通过剪枝掉冗余的空间激活来减少计算量。该方法利用后训练的方式，无需重新训练模型，即可在部署阶段提升效率。通过分析激活图，识别并移除不重要的激活，从而降低后续扫描操作的计算负担。

技术框架：QuarterMap主要包含两个阶段：激活剪枝和维度恢复。首先，对模型的激活进行分析，确定需要剪枝的Token。然后，在四向扫描之前，将这些Token移除。最后，为了保持输出的维度不变，使用最近邻上采样方法将维度恢复到原始大小。整个过程无需重新训练模型，可以即插即用。

关键创新：QuarterMap的关键创新在于其针对状态空间模型的特性，设计了一种高效的后训练Token剪枝方法。与通用的Token合并方法不同，QuarterMap避免了复杂的合并-解合并操作，更加轻量级，更适合SSM的结构。此外，QuarterMap是一种激活剪枝方法，直接作用于激活值，而不是修改模型结构，因此更容易实现和部署。

关键设计：QuarterMap的关键设计包括：1) 激活重要性评估标准：用于确定哪些激活可以被安全地剪枝，这可能涉及到统计激活值的分布、梯度信息等。2) 剪枝比例：决定了需要移除多少比例的激活，需要在吞吐量和精度之间进行权衡。3) 最近邻上采样策略：选择合适的上采样方法，以尽可能减少维度恢复过程中的信息损失。4) 针对不同层或模块，可能需要调整剪枝比例等参数，以达到最佳性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

QuarterMap在ImageNet-1K图像分类任务中，使VMamba的吞吐量提高了高达11%，而精度下降不到0.9%。在ADE20K语义分割任务中，也取得了类似的性能提升。此外，在MedMamba医学影像模型上，QuarterMap在多个医学成像任务中均能提高吞吐量，同时保持精度，证明了其在不同领域和模型上的有效性和泛化能力。

🎯 应用场景

QuarterMap具有广泛的应用前景，尤其是在资源受限的设备上部署视觉状态空间模型。例如，在移动设备上运行VMamba进行图像识别或目标检测，或者在医疗影像分析中使用MedMamba进行疾病诊断。该方法可以显著提高模型的推理速度，降低计算成本，从而加速相关应用的落地。

📄 摘要（原文）

State space models (SSMs) reduce the quadratic complexity of transformers by leveraging linear recurrence. Recently, VMamba has emerged as a strong SSM-based vision backbone, yet remains bottlenecked by spatial redundancy in its four-directional scan. We propose QuarterMap, a post-training activation pruning method that removes redundant spatial activations before scanning and restores dimensions via nearest-neighbor upsampling. Our method improves throughput without retraining. On ImageNet-1K, QuarterMap achieves up to 11% speedup on VMamba with less than 0.9% accuracy drop, and yields similar gains on ADE20K segmentation. Beyond VMamba, we validate QuarterMap on MedMamba, a domain-specific model that shares the same four-directional scanning structure, where it consistently improves throughput while preserving accuracy across multiple medical imaging tasks. Compared to token merging methods like ToMe, QuarterMap is tailored for SSMs and avoids costly merge-unmerge operations. Our method offers a plug-and-play tool for deployment-time efficiency without compromising transferability.