MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods
作者: Zukang Xu, Yuxuan Yue, Xing Hu, Zhihang Yuan, Zixu Jiang, Zhixuan Chen, Jiangyong Yu, Chen Xu, Sifan Zhou, Dawei Yang
分类: cs.LG, cs.AI, cs.CL
发布日期: 2025-01-23 (更新: 2025-03-11)
💡 一句话要点
MambaQuant：提出方差对齐旋转量化方法，实现Mamba模型高效量化
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: Mamba模型 量化 后训练量化 方差对齐 旋转量化
📋 核心要点
- 现有量化方法在Mamba模型上表现不佳，主要由于Mamba结构中的异常值、并行扫描以及通道方差不一致等问题。
- MambaQuant通过KLT增强旋转和平滑融合旋转，自适应地处理通道分布，均衡通道方差，从而实现更有效的量化。
- 实验结果表明，MambaQuant能将Mamba模型的权重和激活量化到8比特，且精度损失小于1%，适用于视觉和语言任务。
📝 摘要（中文）
Mamba是一种高效的序列模型，可与Transformer相媲美，并展现出作为各种任务的基础架构的巨大潜力。量化是神经网络中常用的技术，可以减小模型尺寸并降低计算延迟。然而，将量化应用于Mamba的研究仍然不足，并且现有的量化方法（对CNN和Transformer模型有效）似乎不足以用于Mamba模型（例如，Quarot在Vim-T上即使在W8A8下也会遭受21％的精度下降）。我们率先探索了这个问题，并确定了几个关键挑战。首先，门控投影、输出投影和矩阵乘法中存在明显的异常值。其次，Mamba独特的并行扫描进一步放大了这些异常值，导致数据分布不均匀且呈现重尾分布。第三，即使应用Hadamard变换，权重和激活中各通道的方差仍然不一致。为此，我们提出了MambaQuant，这是一个后训练量化（PTQ）框架，包括：1）Karhunen-Loeve变换（KLT）增强的旋转，使旋转矩阵能够适应不同的通道分布。2）平滑融合旋转，可均衡通道方差，并将其他参数合并到模型权重中。实验表明，MambaQuant可以将权重和激活量化为8位，而基于Mamba的视觉和语言任务的精度损失小于1％。据我们所知，MambaQuant是Mamba系列的首个综合PTQ设计，为进一步推进其应用铺平了道路。
🔬 方法详解
问题定义：论文旨在解决Mamba模型量化的问题。现有的量化方法，例如针对CNN和Transformer设计的量化方法，在Mamba模型上表现不佳，导致显著的精度下降。这是由于Mamba模型中存在异常值、并行扫描以及通道方差不一致等问题，这些问题在传统模型中并不突出。
核心思路：论文的核心思路是通过方差对齐的旋转方法来解决Mamba模型量化中的问题。具体来说，通过Karhunen-Loeve变换（KLT）增强旋转，使旋转矩阵能够适应不同的通道分布，并使用平滑融合旋转来均衡通道方差。这种方法旨在减少量化误差，提高量化后的模型精度。
技术框架：MambaQuant是一个后训练量化（PTQ）框架，主要包含两个核心模块：KLT增强旋转和平滑融合旋转。首先，对权重和激活进行KLT增强旋转，以适应不同的通道分布。然后，应用平滑融合旋转来均衡通道方差，并将额外的参数合并到模型权重中。最后，使用量化器将权重和激活量化到目标比特数（例如8比特）。
关键创新：论文的关键创新在于提出了针对Mamba模型的方差对齐旋转量化方法。与传统的量化方法不同，MambaQuant考虑了Mamba模型特有的数据分布和结构特点，通过自适应的旋转变换来减少量化误差。KLT增强旋转和平滑融合旋转是两个关键的技术创新点，它们分别解决了通道分布不均匀和通道方差不一致的问题。
关键设计：KLT增强旋转的关键在于使用Karhunen-Loeve变换来计算旋转矩阵，该变换能够根据数据的协方差矩阵自适应地调整旋转方向。平滑融合旋转的关键在于设计合适的融合策略，将额外的参数（例如缩放因子）合并到模型权重中，从而减少量化误差。具体的参数设置和损失函数在论文中没有详细描述，属于未知信息。
🖼️ 关键图片
📊 实验亮点
MambaQuant在Mamba模型上实现了显著的量化效果，能够将权重和激活量化到8比特，且精度损失小于1%。实验表明，MambaQuant在基于Mamba的视觉和语言任务上均取得了良好的性能，证明了其有效性和通用性。与现有的量化方法相比，MambaQuant能够更好地适应Mamba模型的特点，从而实现更高的量化精度。
🎯 应用场景
MambaQuant的潜在应用领域包括移动设备、嵌入式系统和边缘计算等资源受限的环境。通过减小模型尺寸和降低计算延迟，MambaQuant可以使Mamba模型在这些平台上高效运行，从而推动Mamba模型在各种实际应用中的部署，例如移动端的自然语言处理、智能监控和自动驾驶等。
📄 摘要（原文）
Mamba is an efficient sequence model that rivals Transformers and demonstrates significant potential as a foundational architecture for various tasks. Quantization is commonly used in neural networks to reduce model size and computational latency. However, applying quantization to Mamba remains underexplored, and existing quantization methods, which have been effective for CNN and Transformer models, appear inadequate for Mamba models (e.g., Quarot suffers a 21% accuracy drop on Vim-T$^\dagger$ even under W8A8). We have pioneered the exploration of this issue and identified several key challenges. First, significant outliers are present in gate projections, output projections, and matrix multiplications. Second, Mamba's unique parallel scan further amplifies these outliers, leading to uneven and heavy-tailed data distributions. Third, even with the application of the Hadamard transform, the variance across channels in weights and activations still remains inconsistent. To these ends, we propose MambaQuant, a post-training quantization (PTQ) framework consisting of: 1) Karhunen-Loeve Transformation (KLT) enhanced rotation, rendering the rotation matrix adaptable to diverse channel distributions. 2) Smooth-Fused rotation, which equalizes channel variances and can merge additional parameters into model weights. Experiments show that MambaQuant can quantize both weights and activations into 8-bit with less than 1% accuracy loss for Mamba-based vision and language tasks. To the best of our knowledge, MambaQuant is the first comprehensive PTQ design for the Mamba family, paving the way for further advancements in its application.