MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods

TL;DR

This work tackles the challenge of post-training quantization for the Mamba family, where traditional Hadamard-based rotations fail to align channel variances due to significant outliers and the PScan amplification. The authors introduce MambaQuant, a PTQ framework combining offline Karhunen-Loève Transformation (KLT)-enhanced rotation with online Smooth-Fused rotation to normalize both maximum values and variances across channels. They demonstrate that their approach preserves accuracy within 1% for 8-bit quantization on Vision and Language tasks and enables 4-bit weights with minimal loss on vision tasks, establishing the first comprehensive PTQ design tailored to Mamba models. The method shows broad applicability across Mamba variants (vision and language) with strong empirical results and manageable memory and computation overhead, offering practical deployment benefits for resource-constrained environments.

Abstract

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.

Figures (13)

• Figure 1: Visualized distribution of hard layers for Mamba quantization. (a) denotes the weight of the gate projection, (b) denotes the input activations of the output projection. (c) represents the output of the parallel scan (PScan) operator, which is also one of the input to the matrix multiplication.
• Figure 2: Maximum values (blue color) and variances (red color) distribution across channels of: (a) the original weight of the gate projection; (b) applying the standard offline Hadamard rotation to (a); (c) applying the proposed KLT-Enhanced rotation to (a); (d) the input activation (generated by PScan) of the matrix multiplication; (e) applying the standard online Hadamard rotation to (d); (f) applying the proposed smooth-fused rotation to (d).
• Figure 3: Mamba block architecture.
• Figure 4: Offline transformation designs utilizing the KLT-Enhanced rotation.
• Figure 5: Fusing smooth parameters into the Mamba structure.