ICLR 2025 | 精度近乎无损！首个Mamba系列模型量化方案MambaQuant-CSDN博客

来源：学姐带你玩AI
本文共1500字，建议阅读5分钟本文提出了一种名为MambaQuant的后训练量化（PTQ）框架，专门针对 Mamba 模型家族的量化问题。

论文题目：

MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods

论文地址：https://arxiv.org/pdf/2501.13484

创新点

作者首次系统分析了 Mamba 模型在量化过程中面临的独特挑战，在门控投影（gate projection）、输出投影（output projection）和矩阵乘法（matmul）层中存在显著的异常值。针对上述问题，作者提出了首个专门针对 Mamba 模型的后训练量化（PTQ）框架 MambaQuant，解决了现有量化方法（如 Hadamard 变换）在 Mamba 模型上效果不佳的问题。
本文通过将 Hadamard 矩阵与 Karhunen-Loève 变换（KLT）矩阵相结合，生成的旋转矩阵能够适应不同通道的数据分布，从而实现方差对齐。在离线模式中，KLT 增强的旋转方法能够有效平衡不同通道的方差，使得量化前的数据分布更加均匀，从而提高量化精度。
本文通过引入平滑参数，对数据进行预处理，使通道间的方差更加均匀。这些平滑参数可以灵活地融入 Mamba 模型的权重中，避免了额外的存储开销。

方法

本文提出了一种名为MambaQuant的后训练量化（PTQ）框架，专门针对 Mamba 模型家族的量化问题。其核心方法是通过KLT 增强的旋转方法和Smooth-Fused 旋转方法来解决 Mamba 模型量化中的关键挑战。KLT 增强的旋转方法通过将 Hadamard 矩阵与 Karhunen-Loève 变换（KLT）矩阵相结合，生成能够适应不同通道分布的旋转矩阵，从而实现方差对齐，使量化前的数据分布更加均匀。Smooth-Fused 旋转方法则通过在 Hadamard 变换之前引入平滑技术，进一步优化通道方差的对齐，并将平滑参数融入模型权重中，避免额外的存储开销。这两种方法共同作用，使得 Mamba 模型在量化后能够保持较高的精度。

Mamba量化中硬层的可视化分布

本图展示了 Mamba 模型中一些难以量化的层的权重和激活值的分布情况。这些层在量化过程中面临显著的挑战，主要原因是它们的数据分布具有显著的异常值（outliers），并且这些异常值会被 Mamba 模型中独特的并行扫描（Parallel Scan, PScan）操作进一步放大。

Mamba 架构

本图展示了 Mamba 模型中不同层在应用不同旋转方法前后的最大值和方差分布情况。这些分布的变化反映了不同量化方法对数据分布的影响，以及 MambaQuant 方法在优化数据分布方面的优势。总体而言，通过对比不同旋转方法的效果，清晰地展示了 MambaQuant 方法在解决 Mamba 模型量化过程中数据分布不均匀问题方面的优势。KLT 增强旋转和平滑融合旋转都能有效地对齐不同通道的方差，使得量化后的数据分布更加均匀，从而提高了量化精度。

利用KLT增强旋转的离线变换设计

本图展示了 MambaQuant 框架中用于离线变换的 KLT 增强旋转（KLT-Enhanced rotation）的设计。这一部分是 MambaQuant 方法的核心内容之一，旨在通过改进传统的 Hadamard 旋转，解决 Mamba 模型在量化过程中面临的方差不一致问题。

实验结果

本表展示了 MambaQuant 在不同视觉任务上的量化性能对比结果。表中列出了多种量化方法（包括 RTN、GPTQ+RTN、SmoothQuant、QuaRot 和 MambaQuant 提出的方法）在不同 Mamba 模型变体（如 Vim-T、Vim-S、Mamba-2d、Mamba-3d 等）上的性能表现，主要关注的是量化精度（accuracy）的损失情况。此外，表中还展示了在不同模型和数据集上的平均性能，进一步证明了 MambaQuant 方法的鲁棒性和有效性。总体而言，Table 1 通过对比不同量化方法在多种 Mamba 模型上的性能，清晰地展示了 MambaQuant 方法的优势，表明其能够在量化权重和激活时保持极高的精度，尤其是在 8 位量化（W8A8）和 4 位量化（W4A8）配置下，显著优于现有的量化方法，是一种非常实用且有效的量化解决方案。

编辑：王菁

关于我们

数据派THU作为数据科学类公众号，背靠清华大学大数据研究中心，分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识，努力建设数据人才聚集平台、打造中国大数据最强集团军。