解决Mamba模型量化难题,实现8比特压缩几乎无精度损失。
MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods
- 用自适应旋转方法处理通道方差不均问题。
- 在视觉与语言任务中实现W8A8量化,精度损失小于1%。
- 适合需要高效部署Mamba模型的研究者与工程师。
Mamba是一种高效的序列模型,性能可媲美Transformer,具有作为通用架构的巨大潜力。量化常用于减少神经网络的模型大小和计算延迟,但将其应用于Mamba仍处于探索阶段。现有针对CNN和Transformer的量化方法对Mamba效果不佳(例如,Quarot在Vim-T$^★$上达到W8A8时准确率下降21%)。我们首次系统研究该问题,发现三大挑战:门控投影、输出投影及矩阵乘法中存在显著异常值;Mamba特有的并行扫描机制加剧了这些异常值,导致数据分布不均且尾部厚重;即使使用Hadamard变换,权重和激活的通道间方差仍不一致。为此,我们提出MambaQuant,一种后训练量化(PTQ)框架,包含:1)增强型Karhunen-Loeve变换(KLT)旋转,使旋转矩阵能适配不同通道分布;2)平滑融合旋转,均衡通道方差,并可将额外参数融入模型权重。实验表明,MambaQuant可在保持低于1%精度损失的前提下,将基于Mamba的视觉与语言任务模型实现权重与激活的8比特量化。据我们所知,MambaQuant是首个针对Mamba家族的完整后训练量化方案,为后续应用铺平道路。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。