Q-MambaIR: Accurate Quantized Mamba for Efficient Image Restoration

📄 arXiv: 2503.21970v2 📥 PDF

作者: Yujie Chen, Haotong Qin, Zhang Zhang, Michelo Magno, Luca Benini, Yawei Li

分类: cs.CV

发布日期: 2025-03-27 (更新: 2025-04-02)

💡 一句话要点

提出Q-MambaIR，用于高效图像恢复的精确量化Mamba模型

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

关键词: 图像恢复 量化 状态空间模型 Mamba 低比特量化

📋 核心要点

  1. 现有SSM在图像恢复中表现出色，但部署到边缘设备受限于内存、算力和功耗，需要高效压缩策略。
  2. Q-MambaIR通过动态调整量化范围的DLS和灵活舍入的RFA，降低量化误差，提升低比特量化下的性能。
  3. 实验表明，Q-MambaIR在图像恢复任务上超越现有量化SSM，实现了更高的精度，同时节省了存储。

📝 摘要（中文）

本文提出Q-MambaIR，一种精确、高效且灵活的量化Mamba模型，用于图像恢复任务。针对状态空间模型（SSM）在低比特量化时性能显著下降的问题，主要原因是异常值加剧了量化误差。Q-MambaIR引入了统计动态平衡可学习标量（DLS），以动态调整量化映射范围，从而减轻由极值引起的峰值截断损失。此外，设计了具有自适应阈值的范围浮动灵活分配器（RFA），用于灵活地舍入值，保留高频细节并维持SSM的特征提取能力。RFA还支持预部署权重量化，平衡了计算效率和模型精度。在图像恢复任务上的大量实验表明，Q-MambaIR始终优于现有的量化SSM，以极小的训练计算量和存储节省实现了更高的SOTA精度。

🔬 方法详解

问题定义：论文旨在解决图像恢复任务中，状态空间模型（SSM）在超低比特（2-4 bits）量化时性能显著下降的问题。现有的量化方法在处理SSM时，由于异常值的存在，会加剧量化误差，导致性能大幅降低。

核心思路：论文的核心思路是通过动态调整量化范围和灵活分配量化值，来减少量化误差，从而提升低比特量化下SSM的性能。具体来说，通过统计动态平衡可学习标量（DLS）动态调整量化映射范围，减轻极值截断损失；通过范围浮动灵活分配器（RFA）灵活舍入值，保留高频细节。

技术框架：Q-MambaIR的整体框架基于Mamba架构，并在量化过程中引入了DLS和RFA两个关键模块。DLS用于动态调整量化范围，RFA用于灵活分配量化值。整个流程包括：输入图像 -> Mamba模块 -> DLS -> RFA -> 输出图像。

关键创新：论文的关键创新在于DLS和RFA的设计。DLS通过统计信息动态调整量化范围，有效缓解了异常值带来的量化误差。RFA则通过自适应阈值进行灵活的舍入，保留了图像的高频细节，维持了SSM的特征提取能力。与现有方法相比，Q-MambaIR能够更有效地处理低比特量化下的异常值问题。

关键设计：DLS的关键设计在于其动态调整量化范围的能力，通过学习一个标量来调整量化映射，从而适应不同的数据分布。RFA的关键设计在于其自适应阈值，该阈值根据数据的统计信息进行调整，从而实现灵活的舍入。此外，RFA还支持预部署权重量化，可以在部署前对权重进行量化，进一步提高计算效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Q-MambaIR在图像恢复任务上显著优于现有的量化SSM模型。例如，在相同的量化比特数下，Q-MambaIR的PSNR指标比其他模型高出多个dB。此外，Q-MambaIR在保持高精度的同时，还实现了显著的存储节省，为边缘设备的部署提供了可能。

🎯 应用场景

Q-MambaIR可应用于各种需要高效图像恢复的场景，如移动设备上的图像增强、低功耗监控系统、以及资源受限的边缘计算设备。该研究的实际价值在于降低了图像恢复模型的计算复杂度和存储需求，使其能够在资源有限的平台上部署，并为未来的低比特量化模型设计提供了新的思路。

📄 摘要（原文）

State-Space Models (SSMs) have attracted considerable attention in Image Restoration (IR) due to their ability to scale linearly sequence length while effectively capturing long-distance dependencies. However, deploying SSMs to edge devices is challenging due to the constraints in memory, computing capacity, and power consumption, underscoring the need for efficient compression strategies. While low-bit quantization is an efficient model compression strategy for reducing size and accelerating IR tasks, SSM suffers substantial performance drops at ultra-low bit-widths (2-4 bits), primarily due to outliers that exacerbate quantization error. To address this challenge, we propose Q-MambaIR, an accurate, efficient, and flexible Quantized Mamba for IR tasks. Specifically, we introduce a Statistical Dynamic-balancing Learnable Scalar (DLS) to dynamically adjust the quantization mapping range, thereby mitigating the peak truncation loss caused by extreme values. Furthermore, we design a Range-floating Flexible Allocator (RFA) with an adaptive threshold to flexibly round values. This approach preserves high-frequency details and maintains the SSM's feature extraction capability. Notably, RFA also enables pre-deployment weight quantization, striking a balance between computational efficiency and model accuracy. Extensive experiments on IR tasks demonstrate that Q-MambaIR consistently outperforms existing quantized SSMs, achieving much higher state-of-the-art (SOTA) accuracy results with only a negligible increase in training computation and storage saving.