点击下方卡片，关注「AI视界引擎」公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
状态空间模型（SsMs）在图像恢复（IR）领域备受关注，因其能够线性扩展序列长度并有效捕捉长距离依赖关系。然而，将SsMs部署到边缘设备面临内存、计算能力和功耗的限制，凸显了高效压缩策略的必要性。
尽管低比特量化是一种通过减小模型尺寸和加速IR任务来高效压缩模型的方法，但SsM在超低比特宽度（2-4比特）下性能大幅下降，这主要是由于异常值加剧了量化误差。为应对这一挑战，作者提出了
MambaIR，一种用于IR任务的精确、高效且灵活的量化Mamba模型。具体而言，作者引入了统计动态平衡可学习标量（DLS），以动态调整量化映射范围，从而减轻极端值引起的峰值截断损失。
此外，作者设计了一种具有自适应阈值的范围浮点灵活分配器（RFA），以灵活地舍入数值。这种方法保留了高频细节并维持了SsM的特征提取能力。值得注意的是，RFA还支持预部署权重量化，在计算效率和模型精度之间取得了平衡。
在IR任务上的大量实验表明，
MambaIR始终优于现有的量化SsMs，在仅略微增加训练计算量和存储节省的情况下，实现了远超当前最优（SOTA）的精度结果。
unsetunset1. 引言unsetunset
图像修复（IR）自数十年前以来已被广泛认为是计算机视觉和图像处理领域的一个重要问题。它专注于从因噪声、模糊或压缩伪影导致的退化或损坏版本中恢复原始图像。当前最先进的方法基于深度神经网络，该网络结合了卷积神经网络（CNNs）[15, 19]进行局部特征提取和Transformer[18, 33]对图像内的长距离依赖关系进行建模。
最近，状态空间模型（SSMs）[9, 10]被提出用于以相对于序列长度的线性缩放计算来建模长距离依赖关系。然而，基于SSM的架构的巨大模型规模仍然是实时部署在资源受限设备上的挑战。这促使IR领域需要SSM压缩方法，而这一方向仍基本未被探索。
低比特量化是压缩深度学习（DL）模型和加速模型推理的有效方法。通过将浮点精度的权重和激活值转换为低比特整数值（4比特或更少），可以显著降低模型内存占用和延迟。然而，对状态空间模型（SSMs）进行量化存在重大挑战，在低比特设置下会导致严重的量化失效和性能退化。
首先，系统中的结构异常值持续存在（见图3），它们与输入数据的固有变异性相互作用，为静态量化带来了基本挑战。具体而言，为了适应异常值，量化范围必须扩展。然而，这种扩展不可避免地降低了数据密集区域（大多数值集中于此）的表示密度。相反，如果为了保持密度而限制范围，则可能截断异常值的关键信息。这种权衡源于传统量化范围的静态设计。这些范围是预先设定的，对所有输入实例保持固定，缺乏适应动态数据分布的能力。其次，低比特量化中普遍存在的确定性舍入导致反向梯度不匹配和正向参数同质化。这会导致基本的离散化误差，引发非线性误差传播甚至特征退化。由于离散化，这迫使在保留精度和实现高计算效率之间做出次优妥协。在超分辨率任务中，同时保留高频纹理和Sparse极值激活至关重要。
受上述见解启发，作者首先提出一种动态平衡可学习标量（DLS）。该方案通过设计可学习标量，通过捕获激活分布特征动态适应量化范围。作者进一步提出一种范围浮点灵活分配器（RFA），这是一种使用自适应阈值对值进行舍入的灵活离散化机制，具有灵活的分配器。RFA已被证明通过局部梯度保持有效缓解参数同质化，从而保持SSM模块的特征表示能力。
大量实验表明，作者的QMambaIR在各种位宽下均优于当前最先进（SOTA）的量化方法，包括训练后量化（PTQ）和量化感知训练（QAT）方法（参见表1及3至5），在图像恢复（IR）模型方面。值得注意的是，作者的4位Q-MambaIR在性能上与使用8位的现有PTQ方法相当。实验结果表明，所提出的QAT方法显著提升了模型精度，即使采用激进的低2位权重量化的QAT也是如此。作者还通过与QAT和PTQ量化网络的对比，证明了作者
MambaIR的有效性，突显了其相对于不同方法的优越性。作者的主要贡献总结如下：
作者提出了Q-MambaIR，一种用于高效图像恢复的精确Mamba模型。它缩小了低比特量化模型与其全精对应模型之间的差距，使得在资源受限场景下的有效部署成为可能。为了解决极端值保留问题，作者提出了动态平衡可学习标量（DLS）。具体而言，作者的SDLQ通过可学习感知器调整量化范围，通过捕捉激活分布的统计特征，有效减少截断损失。为了增强高频特征保留，作者提出了范围浮动灵活分配器（RFA），它通过自适应阈值离散化保留关键细节，同时确保强大的特征提取能力，从而提高图像恢复任务中的重建质量。
作者采用Q-MambaIR将基于SSM的IR网络压缩至更低的位宽，同时保持具有竞争力的恢复质量（如图2所示），从而得到相应的量化 Baseline 。Q-MambaIR在SOTA量化SSM方法上实现了更优的性能。
unsetunset2. 相关工作unsetunset
用于图像修复的Mamba。近年来，卷积神经网络（CNN）的集成从根本上改变了传统的图像修复方法，通过利用神经网络的固有优势显著提升了性能[15, 19]。随后，基于Transformer的方法通过有效建模长距离依赖关系并捕捉空间上遥远像素之间的复杂关系，进一步推动了该领域的发展。最近，MambaIR[10]因其创新性地应用状态空间模型（SSM）线性递归特性而受到研究界的广泛关注，这些特性能够实现高效的长距离依赖建模。然而，它存在模型尺寸庞大和计算需求密集的问题，这严重阻碍了其在资源受限环境下进行高效图像处理的部署。在这种情况下，低比特量化作为一种吸引人的策略应运而生，它有望在保持计算效率的同时，利用SSM的优势进行图像修复任务。
IR/Mamba量化。量化方法根据是否涉及训练可分为两类：量化感知训练（QAT）和训练后量化（PTQ）。QAT涉及使用量化参数重新训练神经网络，使模型能够适应量化引起的扰动，并实现改进损失的最优解[3, 6, 8]。相比之下，PTQ直接将量化应用于预训练模型，无需额外训练，提供更高效的计算实现，尽管可能以降低精度为代价。
SSM的线性递归特性为量化引入了独特的挑战，尤其是在激活量化方面。异常值的存在对量化构成挑战。当前针对SSM的PTQ方法主要通过基于旋转的技术来解决这些挑战。例如，Quamba提出了一种静态的每张量8位量化方法，结合Hadamard变换来平滑激活值[5]。类似地，MambaQuant采用KLT增强旋转和Smooth-Fused旋转来缓解激活量化问题[25]。然而，这些方法目前仅限于8位量化，并且主要应用于语言模型领域。
量化感知训练（QAT）凭借其基于训练的优化方法，在实现超低比特量化方面具有显著潜力。例如，QuantSR [23] 展示了在低比特量化下进行精确高效的超分辨率（SR）处理的可行性 [11, 16, 23]。然而，传统方法受限于其依赖静态量化范围，这限制了它们在处理具有易出现异常激活的SparseSparse映射（SSM）时的有效性。因此，迫切需要进一步研究开发针对SSM的量化方法，以实现IR中鲁棒且高效的超低比特量化。
unsetunset3. 方法unsetunset
为解决SSM低比特量化问题，作者首先介绍了MambaIR中使用的核心部分——视觉状态空间模型（VSSM）的概述，随后分析了现有低比特量化架构的固有局限性。基于此分析，作者提出了Q-MambaIR来解决这些问题，这是一种用于高效推理的精确量化Mamba（概述如图4所示）。Q-MambaIR主要由两种技术组成：动态平衡可学习标量（DLS）和范围浮点灵活分配器（RFA）。DLS通过动态适应激活分布，克服了确定性量化范围的限制，从而最小化由异常值裁剪引起的截断损失。RFA引入了一种具有灵活舍入阈值的软舍入机制，减轻了静态舍入方案通常累积的错误。
3.1. 前提
视觉状态空间模型。VSSM模块是MambaIR中的主要功能特征提取器VSSM的架构如图5(a)所示。它在不牺牲其优势的情况下将SSM应用于视觉数据。具体来说，SSM通过隐藏状态
定义了一个序列到序列的转换
。它包含两个阶段，并通过状态空间方程捕获长距离依赖关系：
快速推理的离散化阶段：
高效并行化训练的计算阶段：
VSsM量化。作为一种常见的量化方法，量化器将计算单元中使用的权重和激活映射到均匀的离散序列：
其中
表示低比特宽度的量化权重或激活值，
是浮点格式的输入张量。函数
确保在舍入前值被约束在
-位整数的可表示范围内。
是缩放因子，它决定了
的范围如何映射到
和
之间的舍入范围。
是位移因子，用于调整零点对齐，影响量化范围分布。
局限性。现有低比特量化架构在应用于图像恢复任务中的SSM时面临显著局限性。(1) 虽然alpha和beta等可学习参数能够实现每个算子 Level 的动态量化范围调整，但它们仍然为每个算子提供固定的量化范围，缺乏对输入图像变化的适应性。这种僵化性导致大量信息丢失，特别是在存在异常激活值的情况下，而异常激活值在SSM中很常见。(2) 传统的静态舍入机制在反向传播过程中引入离散化误差和梯度不匹配，这会降低恢复图像的精细细节。这些误差在量化过程中累积，最终损害SSM精确表示复杂细节的能力，而细节是高质量图像恢复的关键方面。这些局限性突显了需要针对视觉任务中SSM独特特性的更具适应性和灵活性的量化策略。
3.2. 动态平衡可学习标量 (DLS)
为解决现有方法的第一个局限性，作者为Q-MambaIR提出了一种动态平衡可学习标量（DLS）。通过动态量化范围，DLS通过根据实时输入统计数据自适应调整缩放因子和范围边界来应对这一挑战。
DLS采用一种简单的线性投影方法，将复合向量映射到缩放因子
，仅需八个参数。它有效地捕捉了输入数据的统计分布特征，建立了最优量化范围与数据固有分布特征之间的投影关系。缩放因子
和偏移因子
被定义为复合特征向量的线性加权总和，如下所示：
α
β
此处，复合特征向量被定义为输入的四个统计特征的组合，即
。3
和
是可学习的权重向量。
作者提出的
量化器被嵌入在整个量化后的SS2D块中（如图5中橙色箭头所示），用于量化感知训练，以将激活值压缩到低位宽。
通过根据实时输入统计数据自适应调整量化范围，DLS解决了两个关键问题：固定范围与动态数据分布的不兼容性，以及异常值覆盖与表示密度之间的固有权衡。如图3所示，在SSM结构中观察到异常值，Quamba [5] 在语言任务中以及manbaquant [25] 也注意到了这一现象。为解决此问题，DLS的核心机制采用双阶段自适应。当检测到结构异常值时，DLS暂时扩展量化范围以防止信息截断。在正常条件下无异常值时，它收缩范围以在数据密集区域保持高分辨率表示。这种动态行为由轻量级统计学习模块控制，这些模块持续监控输入特征。
因此，该方法不仅减少了由异常值裁剪引起的精度损失，还优化了不同场景下的比特分配效率。
3.3. 范围浮动灵活分配器 (RFA)
为解决现有方法的第二个局限性，作者提出了一种范围浮动灵活分配器（RFA）。与传统的固定舍入方法不同，RFA采用了一种阈值自适应机制，根据张量统计数据动态调整舍入边界，并结合梯度近似技术来使前向传播和反向传播保持一致（如图4所示）。它由软前向阈值（SFT）和软反向近似（SBA）组成。SFT将输入
映射到最近的量化 Level
。
SBA通过分段线性插值构建可微分的梯度路径。固定斜率（0.1）用于近似圆Sa函数梯度
的离散量化 Level 。其中，
表示量化步长，构建均匀分布的整数区间。可学习的量化阈值参数
与网络权重联合优化，使量化区间能动态匹配SSM激活值的分布。范围浮动机制通过灵活分配量化阈值至关键区域（如SSM权重中的高频分量），减少离散化误差，从而在纹理提取过程中减轻细节损失。重要的是，量化权重严格遵循均匀区间，确保与标准定点运算单元（如INT8 Tensor Cores）直接兼容，因此在推理过程中引入零运行时开销。因此，这种自适应舍入与梯度优化的协同设计显著实现了高频纹理的低损失压缩，同时保持推理效率。
unsetunset4. 实验unsetunset
4.1. 设置
4.1.1. 数据集与评估
遵循先前研究[10]的设置，作者在多种图像处理任务上开展实验，包括图像超分辨率（即经典超分辨率和轻量级超分辨率）和图像去噪（即高斯彩色图像去噪）。经典超分辨率训练使用DIV2K[24]和Flickr2K[19]数据集的组合，而轻量级超分辨率模型则仅在DIV2K[24]数据集上进行训练。超分辨率方法的基准测试采用五个标准数据集：Set5[2]、Set14[28]、B100[21]、Urban100[12]和Manga109[22]。高斯去噪模型使用DIV2K[24]、Flickr2K[19]、BSD500[1]和WED[20]进行训练，评估则在BSD68[21]、Kodak24、McMaster[31]和Urban100[12]上进行。定量分析采用在YCbCr空间的亮度（Y）分量上计算得到的PSNR和SSIM指标。此外，实验还涉及JPEG压缩伪影的减少。
4.1.2. 提出的量化 Baseline
作者使用所yinqing-Q-MambaIR_2503对MambaIR模型的权重和激活值进行了量化。在所有对比模型中，作者使用低比特宽（例如2、4或8）对网络的主体部分进行量化。为了进行公平的比较，对比方法的实现遵循官方代码，并使用与作者相同的设置进行训练。对于PTQ方法，作者遵循Mamba量化策略Mambaquant [25] 和Quamba [5]。对于QAT方法，作者使用QuantSR [23] 和LSQ [3] 对SSM进行量化。作者将
-比特权重和
-比特激活值表示为
。第4.2节和第4.3节展示了完整的效率性能。
4.1.3. 训练策略
根据先前研究[4, 10]，作者通过应用水平翻转以及随机的
、
和
旋转进行数据增强。此外，在训练过程中，作者将原始图像裁剪为
的块用于图像超分辨率（SuperResolution）和
的块用于图像去噪。模型训练20k次，每个训练批次包含8个图像块用于图像超分辨率和4个用于图像去噪。为优化模型，作者采用Adam优化器。学习率初始设置为轻量级超分辨率为
，其他任务为
，并在训练迭代达到特定里程碑时减半。作者的Q-MambaIR模型使用2块NVIDIA A100 GPU进行训练。
4.2 图像超分辨率对比
作者选择MambaIR作为作者的 Backbone 网络，并将yinqing-Q-MambaIR_2503与最近期的量化方法进行比较。作者选择了两种Mamba量化PTQ方法，即Mambaquant [10]和Quamba [5]。对于QAT方法，作者评估了QuantSR [23]和LSQ [8]的性能。
经典图像超分辨率。表1展示了不同方法在经典图像SR任务上的比较。作者的4位QMamba在尺度为
时，其PSNR/SSIM得分与8位Quamba和Mambaquant相当或更优。在4位情况下，作者的QMamba在Urban100数据集上比LSQ实现了0.21 dB/0.0014
和
更高的PSNR/SSIM值。在2位设置下，作者的QMamba在Urban100数据集上比QuantSR显著提升了
和1.37 dB/0.0522
。在更低的位宽下，作者的QMamba相比先前技术有显著改进，并缩小了低位模型与全精度模型之间的差距。
轻量级图像超分辨率。作者使用yinqing-Q-MambaIR_2503量化Mamba-light模型，并在表5中展示了其性能。在4比特的情况下，作者的
MambaIR在
尺度Urban100数据集上比LSQ [8]高0.21dB PSNR。在2比特的情况下，作者的QMambaIR在
尺度Urban100数据集上比LSQ [8]高2.49dB PSNR。yinqing-Q-MambaIR_2503为超低比特量化提供了一种可能的解决方案。
模型复杂度。通过超低比特量化，量化模型可以实现非常显著的压缩比。在表2中，作者分别展示了基于参数量和运量的压缩比和加速比。遵循QuantSR [23]，作者仅对高层特征提取器中的权重和激活进行量化，但作者计算了整个模型的压缩比和加速比。在2比特设置下，作者的Q-MambaIR在轻量级MambaIR中实现了约90.6%和88.0%的参数量和运算量压缩比。
4.3. 其他IR任务的比较
高斯彩色图像去噪。高斯彩色图像去噪的结果如表3所示。作者选取了噪声水平为25的典型案例。可以看出，yinqing-Q-MambaIR_2503在性能上优于所有对比方法。具体而言，在Urban100数据集上，使用4位量化时，yinqing-Q-MambaIR_2503比QuantSR [23]方法提高了高达0.3 dB。JPEG压缩伪影减少。作者进一步评估了Q-MambaIR在JPEG压缩伪影减少方面的鲁棒性，质量因子为30。结果如表4所示。结果表明，在SIDD数据集上，yinqing-Q-MambaIR_2503在2位量化时PSNR达到了0.52 dB，优于现有的量化方法。
4.4. 消融研究
为验证作者Q-MambaIR所采用技术的有效性，作者在DLS和R上进行了消融实验。作者采用MambaIR作为图像恢复 Backbone 网络，并使用4位和2位量化设置进行20K次迭代训练。为建立量化 Baseline ，作者使用了原始量化方法QuantSR [23]。随后，作者将DLS和/或RFA整合到MambaIR中并将其二值化。实验结果汇总于表6，其中作者报告了在五个测试基准上的PSNR值。
Q-MambaIR组件的消融实验。基础 Backbone 网络展示了低比特量化的可行性。作者提出的RFA显著提升了量化范围分配的灵活性，提高了量化MambaIR的性能。如表6所示，仅加入DLS相较于 Baseline 在Urban100上精度提升了1.05 dB。然而单独使用RFA无法完全恢复SSM中的截断损失。当与RFA结合时，yinqing-Q-MambaIR_2503实现了更优的细节保留，达到PSNR为32.43 dB。这缩小了与对应实数MambaIR的精度差距至仅1.5%，突显了所提方法的有效性。
RFA初始化策略的消融实验。作者进一步研究了DLS中量化范围初始化的各种设计。具体而言，作者将缩放因子初始化为
和
，同时将偏移因子设置为输入的平均值
以及最小值和最大值的中点
。这些参数被组合成四个控制组进行实验。表7中的结果表明，由
和
组成的控制组取得了最佳性能，验证了作者初始化策略的有效性。
4.5. 可视化
视觉效果。在图1中，作者展示了在ClassicSR任务上具有代表性的几种方法在2比特和4比特情况下的视觉效果，比例尺为
。对于每种情况，作者将其与几种量化方法进行比较，包括LSQ [8]、QuantSR [23]、Quamba [5]和Mambaquant [25]。与其他量化方法相比，QMambaIR能够恢复更多的结构细节并减轻更多的模糊伪影。因此，作者的Q-MambaIR与全精度MambaIR [10]之间的视觉差异较小。这些视觉比较进一步证明了作者Q-MambaIR的有效性，这与表1中的观察结果一致。
unsetunset5. 结论unsetunset
在本文中，作者提出了Q-MambaIR，一种用于图像恢复任务的精确、高效且灵活的量化Mamba网络，解决了基于SSM模型的低比特量化中的关键挑战。所提出的网络采用统计动态平衡可学习标量（DLS）进行动态量化范围调整，并使用范围浮动灵活分配器（RFA）保留细节梯度。
作者有效缓解了异常值适应、梯度不匹配和参数同质化等问题，同时保留了高频纹理。通过大量实验证明，QMambaIR在准确性和计算效率方面均优于现有的最先进的量化图像恢复网络。
点击上方卡片，关注「AI视界引擎」公众号