MambaQuant：基于方差对齐旋转方法的 Mamba 家族量化
MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods
AI总结：
本文提出 MambaQuant 训练后量化框架，通过 KLT 增强旋转和 Smooth-Fused 旋转解决 Mamba 模型中的离群值与通道方差不一致问题，首次实现 Mamba 家族权重和激活 8 位量化且准确率损失低于 1%。
摘要：
Mamba 是一种高效的序列模型，可与 Transformer 相媲美，并展现出作为多种任务基础架构的巨大潜力。量化通常用于神经网络中以减小模型规模和计算延迟。然而，将量化应用于 Mamba 的研究仍处于探索不足的状态，并且现有的量化方法虽然在 CNN 和 Transformer 模型上有效，但对 Mamba 模型似乎并不适用（例如，Quarot 在 W8A8 设置下于 Vim-T$^\dagger$ 上仍出现 21% 的准确率下降）。我们率先对该问题进行了探索，并识别出若干关键挑战。首先，门控投影、输出投影和矩阵乘法中存在显著的离群值。其次，Mamba 独特的并行扫描进一步放大了这些离群值，导致数据分布不均匀且呈重尾分布。第三，即使应用 Hadamard 变换，权重和激活中跨通道的方差仍然不一致。为此，我们提出了 MambaQuant，一个训练后量化（PTQ）框架，包括：1）Karhunen-Loeve 变换（KLT）增强旋转，使旋转矩阵能够适应不同的通道分布。2）Smooth-Fused 旋转，用于均衡通道方差，并可将额外参数合并到模型权重中。实验表明，对于基于 Mamba 的视觉和语言任务，MambaQuant 可以将权重和激活均量化为 8 位，且准确率损失低于 1%。据我们所知，MambaQuant 是首个面向 Mamba 家族的全面 PTQ 设计，为其应用的进一步发展铺平了道路。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
