MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods

We sincerely thank you for your valuable time and efforts in reviewing our manuscript, and we tried our best to address all your concerns as follows.

The KLT needs more empirical evidence on more diverse data.

Thanks for this constructive comment. To evaluate the robustness of our KLT method, we perform the following experiments.

1. In Table 2 of the manuscript, we evaluate our MambaQuant on Mamba-LLM. The calibration data is sampled from the HellaSwag dataset. Zero-shot evaluation results, especially on other datasets like ARC-E, ARC-C, PIQA, and Winogrande, can effectively demonstrate the generalization ability of our KLT method.

2. We have added visualized examples in Figure 12 of Appendix A.7. It can be clearly seen that: For inputs outside the calibration set and with similar distribution, our KLT-enhanced rotation can also achieve significant optimization; For inputs outside the calibration set but with distinct distribution, the effectiveness of the KLT-enhanced rotation is still not worse than the Hadamard rotation.

3. It should be mentioned that it is a common practice for PTQ to characterize a wide range of data distribution with a static subset, which works effectively in most cases.

The methods proposed are only studied in the context of, and only apply to, the Mamba family of architectures. It would be great if the methods here transferred to different architectures.

First of all, this work focuses on quantization for the Mamba family models.

Second, our methods also can be well generalized to other families such as the mainstream Transformer-based Large Language Models (T-LLMs). The inconsistent variance problem of the Hadamard rotation arises from the algorithm itself and does not depend on any specific model structure. Despite the imperceptible effect on T-LLMs, this limitation is particularly significant when quantizing the Mamba models.

1. For the KLT-enhanced rotation, it essentially changes the rotation matrix from $H$ to $KH$ as stated in the Equation (11). Considering that the rotation scheme is quite mature for T-LLMs [1,2], our KLT-enhanced rotation can be directly applied.

2. For the Smooth-fused rotation, it uses the smoothing technique to uniform the channel variances before the online rotation. We further integrate the extra smoothing factors into weights with consideration of the Mamba structure. Since both the smoothing methods [3,4] and rotation methods are well-developed for T-LLMs, our smooth-fused rotation can also be easily applied.

Lastly, we further explore the robustness of the algorithms in MambaQuant on T-LLMs. As shown in the table below, we quantize llama2-7b to w4a4 using the algorithms of Quarot and MambaQuant respectively. It can be clearly concluded that the KLT-enhanced rotation also outperforms the SOTA baseline.

Is there any aspect of runtime which is significantly improved or made worse by these changes? It would be interesting to see how these bear out in practice.

Thanks for this constructive suggestion.We compared the average inference speeds before and after quantization based on the Mamba-2.8b model with seqlen 512 in A6000. Under w8a8 quantization setting, the acceleration ratio is 1.2 times. Due to the limited rebuttle time, our kernel has not been highly optimized, and there is still room for improvement in speed.

If there are still any unresolved doubts, please feel free to let us know, and we will make every effort to solve them.

References:

[1] Quarot: Outlier-free 4-bit inference in rotated llms.

[2] SpinQuant--LLM quantization with learned rotations.

[3] Smoothquant: Accurate and efficient post-training quantization for large language models.

[4] Omniquant: Omnidirectionally calibrated quantization for large language models.

Dear Reviewer,

Thank you for your constructive comments and valuable suggestions.

We are glad to address all your questions and sincerely appreciate your recognition.

Best, Authors

We sincerely thank you for your valuable time and efforts in reviewing our manuscript, and we tried our best to address all your concerns as follows.

Could you elaborate on any additional differences between Mamba-PTQ and your work?

When we were writing our paper, we surveyed mainstream AI conferences and journals but did not find Mamba-PTQ [1]. Up to now, we note that this preprint has not yet passed the review of professional peers.

Based on your suggestion, we have carefully reviewed Mamba-PTQ. The most closely related aspect is its discussion of outliers in Mamba networks. While both works identify the existence of outliers, our study goes a step further by revealing that the unique PScan operator in Mamba networks can amplify these outliers. More importantly, our paper offers several unique contributions, which we outline below:

1. Variance viewpoint. We analyze not only the issue of outliers, but also the importance of the consistent channel variances for quantization. Balanced variance is essential for reducing quantization error but typically overlooked by existing works.

2. KLT-enhanced rotation. We theoretically analyze why existing advanced rotation[4,5] techniques fail to address the issue of variance inconsistency (see details in Section 4.1) and propose a variance-consistent KLT-enhanced rotation technique (see details in Section 4.2).

3. Smooth-Fused scheme. As metioned In Section 4.3, our Smooth-fused scheme integrates Smooth parameters into model weights during the offline phase, avoiding any additional computation during online inference. In contrast, Mamba-PTQ adopts traditional online smoothing which introduces additional computational overhead.

4. MatMul quantization. Mamba-PTQ only quantizes linear layers, while our work further quantizes the MatMul operator. This operation is hard to be quantized as effected by the PScan and also consumes large inference latency especially for the prefill phase.

5. Superior performance. Our work achieves nearly lossless quantization accuracy (less than 1%) in W8A8 setting on both computer vision and large language tasks. In contrast, Mamba-PTQ does not report its performance on computer vision tasks, and its quantization incurs significant accuracy loss on large language tasks. For instance, the quantization accuracy of Mamba-PTQ in W8A8 drops by more than 10%, with the Mamba-790 model experiencing a drop of over 20%.

It seems that KLT enhanced rotation is very straightforward. Is there is any additional insight to it?

Obtaining the matrix K is not theoretically complex as described in your comment (eigenvalue decomposition on the covariance matrix). Although, we would like to stress the motivation and the effectiveness of this approach besides its computational process.

1. Through detailed analysis in Section 4.1, we reveal the shortcoming of the current rotation methods. They do not take the matter of consistent variances into consideration, leading to sub-optimal performance for the Mamba family. We thus introduce the method of KLT, which has not been adopted before to solve this problem.

2. To provide readers with a clearer insight into the effectiveness of our KLT method, we present Equation (8) with (12), and Equation (9) with (13) in a comparative way. As shown in Equation (13) of the manuscript, this method is simple but can effectively equalize the channel variances.

3. The rotation scheme [4,5] has not been applied to the Mamba structures before. We effectively apply it to the Mamba structure.

In conclusion, the KLT method is practically simple and also significantly effective for quantization, promoting the deployment of the Mamba family models. Our contribution is explaining why it works and how it works.

Would be great if authors could include detailed per-task results for LM-eval benchmark in the appendix, to improve reproducibility and make it easier to compare to.

We have included all per-task results for LM-eval benchmark in Table 6 in the Appendix (A.6). We also plan to open source all related codes and data to promote the development of the community.

If there are still any unresolved doubts, please feel free to let us know, and we will make every effort to solve them.

References:

[1] Mamba-PTQ: Outlier Channels in Recurrent Large Language Models.

[2] Smoothquant: Accurate and efficient post-training quantization for large language models.

[3] Omniquant: Omnidirectionally calibrated quantization for large language models.

[4] Quarot: Outlier-free 4-bit inference in rotated llms.

[5] SpinQuant--LLM quantization with learned rotations.

We sincerely thank you for your valuable time and efforts in reviewing our manuscript, and we tried our best to address all your concerns as follows.

This paper does not detail how to get K for different input X.

As stated in line 263 in Section 4.2 of the manuscript, K is generated from calibration data. We have also added explicit statements in the Section 4.2 of the manuscript to stress this point. This statically generated $K$ can be well-generalized for different input $X$ .

Sure, here is a detailed example showing how to obtain a single $K$ from the calibration dataset.

Step 1:

Supposing that the shape of each activation sample $X_i$ is (L, D) (where L is the sequence length and D is the hidden dimension), the calibration dataset [$X_1$, $X_2$, $X_3$, ..., $X_N$] can be seen as a tensor with shape (N * L, D). This tensor corresponds to the character $X$ in the Equation (10) of the manuscript.

Step 2:

The Equation (10) first calculates the covariance matrix ${C}_X$ of $X\in \mathbb{R}^{NL,D}$, which can be written as: ${C}_X=\frac{1}{NL-1}{X}^T{X},$ where the covariance matrix $C_X$ has shape (D, D).

Step 3:

The Equation (10) then performs eigenvalue decomposition (ED) to $C_X$: $\text{ED}(C_X)=\frac{1}{NL-1}{K\Lambda} {K}^T.$ Thereby, a single $K$ with shape (D, D) that can be obtained from the calibration dataset with N samples. This $K$ is actually the eigenvector matrix of the covariance matrix $C_X$ of the calibration data $X$.

More explanation:

The essence of the KLT method is to perform ED to the covariance matrix. The tensor shape of the covariance matrix is irrelevant to the number of samples (i.e., it is always (D, D)), owing to the calculation of $X^T X$. Thereby, we can always generate a single $K$ from the calibration dataset.

We sincerely thank you for your valuable time and efforts in reviewing our manuscript, and we tried our best to address all your concerns as follows.

I think it is possible (also easier) to control the variance of Hadamard transformed weight by simply applying a diagonal scaling matrix. How does this scheme compare to the proposed method.

Our KLT-enhanced rotation technique stands out by ensuring both consistent variances and uniform maximum values concurrently.

(1) Constrains of the simple scheme.

A diagonal scaling matrix does help in aligning variances. However, neither applying it before nor after the Hadamard rotation can simultaneously ensure consistent variances and uniform maximum values. Given input data $X$, diagonal scaling matrix $S$ with variances as values, and the Hadamard matrix $H$, we can establish three transformed rotation matrices: $SH$, $HS$, and $KH$ (ours).

1. For $SH$ rotation, through multiplying $X$ with $S$, the variances can be temporarily aligned. However, this occurrence will be immediately disturbed by multiplying $H$, as Hadamard rotation can change the values of variances.
2. For $HS$ rotation, through multiplying $X$ with $H$, the maximum values can be temporarily uniformed. However, this occurrence will be immediately disturbed by multiplying $S$, as the value difference of $S$ can lead into new outliers.

(2) Effectiveness of our approach.

We propose KLT-enhanced rotation, which overcome the above-mentioned constrains. Considering the Equation (9) in the manuscript, $H_K=KH,$ and the rotation to data $X$ can be described as: $XH_K=XKH.$ Through multiplying $X$ with $K$, the distribution of $X$ is transformed but not variances-aligned yet. By multiplying $H$, the $XK$ is endowed with uniformed maximum values (owing to the property of Hadamard transformation [1]). At the same time, due to the transformed distribution of $X$, the final channel variances are also aligned (as detailedly proved in Section 4.2 of the manuscript).

(3) Further proof via experimental results.

Following the experimental settings in Section 5 of the manuscript, we quantize the Vision Mamba models [2] in W8A8 setting by using different rotation matrices. Results show that our KLT-enhanced rotation (KH) outperforms other simple schemes.