Q-Mamba: Towards more efficient Mamba models via Post-Training Quantization

Q1: Outliers are a well-established challenge in neural network quantization, not just in language models but across neural architectures. It is unclear the unique challenges that SSMs pose over standard transformer architectures. While the authors claim this is the first paper to study SSM quantization, it is unclear where the unique challenges arise. Enough prior studies have investigated fine-grained scaling factors and block-wise reconstruction to warrant benchmarking against those techniques. Ultimately, it is unclear if existing techniques are already sufficient to solve the challenges the authors present. For example, many PTQ works have investigated groupwise quantization (e.g., GPTQ) to increase the granularity. This also establishes a grid of scaling factors similar to the decoupled scale quantization proposal. Finally, rotation-based quantization techniques have been studied in a handful of papers over the last year (e.g., QuIP, QuaRot, SpinQuant). There is no mention of any of these approaches in the paper, and it has been one of the most effective techniques to reducing outliers in language models.

R1:In this study, we emphasize that Q-Mamba focuses on the quantization of state caches in Mamba architectures, setting it apart from previous research on quantizing transformers and CNNs. States in Mamba are expanded to be $N$-times larger than standard activations, where $N$ represents the state dimension (128 in Mamba-2 models). Section 4 demonstrates that these state caches contribute significantly to memory consumption, particularly after weights are quantized to low-bit representations. It is important to note that the quantization of linear projections for Mamba models is not the primary contribution of this work, although we present preliminary experiments on Mamba-1 and Mamba-2 using GPTQ and SmoothQuant. On the other hand, Q-Mamba is orthogonal and can be combined with previous quantization methods for linear layers, such as OmniQuant, AWQ, and Mamba-PTQ.

Q2: It is unclear why BRECQ (or other existing block-wise reconstruction algorithms) can't be applied to SSMs. The analysis claims parallel training and layer selection are the challenges, but it is not clear if (or how) this is a restriction on the algorithm or some implementation of it. Can you please provide more detail?

R2: We emphasize that adjusting parameters for quantized states is a non-trivial task in block-wise reconstruction. SSMs (and other linear attention mechanisms) rely on the equivalence between two computational modes: the quadratic mode for training and the recurrent mode for inference. This equivalence depends on the linear nature of the hidden-to-hidden transformations, $h_t = A_t h_{t-1} + B_t x_t$. In contrast, RNNs, defined by $h_t = \sigma(A_t h_{t-1} + B_t x_t)$ (where $\sigma$ is a nonlinear activation function), cannot be parallelized during training. To minimize memory bandwidth utilization, we store state caches as low-bit elements and subsequently load and dequantize them before computation at the next timestep. This process introduces a new sequence transformation, defined as $h_t^q = \bar{A} Q(h_{t-1}^q) + \bar{B} x_t$, where $Q(\cdot)$ is a non-linear quantization function. It is crucial to distinguish $h_t^q = \bar{A} Q(h_{t-1}^q) + \bar{B} x_t$ from the quantized value of the original $h_t$, i.e., $Q(h_t) = Q(\bar{A} h_{t-1} + \bar{B} x_t)$. A naive approach would directly apply recurrent mode for token-by-token generation in the training. However, this method becomes computationally prohibitive for large input lengths (e.g., 2048). To address this, we propose Efficient Selectivity Reconstruction (ESR), which utilizes the parallel algorithm to compute $h_t$ across all timesteps and simulates quantization errors by quantizing only a single step during training.

Q3: The notation in Equation 1 confuses me. The matrix dimensions don't seem to line up. Should there be a transpose for $A_t h_{t-1}$ since $A_t$ is $N \times N$ and $h_{t-1}$ is $T \times N$? Also, I imagine $B_t x_t$ is an outer product and would also need to be transposed since $x_t$ is $T \times P$ and $B_t$ is $N \times 1$. If this understanding is correct, can you please explain where the $N \times P$ comes in?

**R3:**We apologize for the confusion. We have significantly modified the background section.

Q1: Computation Overhead in Quantization. The quantization and dequantization process requires computing an outer product of two vectors, which can be computationally expensive compared to other group quantization methods available in the field.

R1: During the generation stage, state caches $h_t$ are quantized at each time step to reduce memory consumption and are subsequently dequantized in the next step, which is similar to KV cache quantization. We implement a cuda kernel that load quantized $h_t{-1}$, compute SSM $h_t=A_th_{t-1}+B_tx_t$, then quantize and pack $h_t$ and store into memory. The computation overheads for scales are neligiable compared to original SSMs computation and reduce half memory-bandwith utlizaition.

Q2:Lack of Comprehensive Evaluation: The experimental results lack sufficient diversity across tasks and model sizes. Evaluating Q-Mamba’s performance on additional Mamba family models would strengthen its claims of versatility.

The results of different tasks (including generation tasks and zerot-shot tasks) and model sizes (from 130M to 2.7B) are shown in Table 1 and Table 2.

Q3:Limited Baseline Comparisons: The paper could benefit from comparisons with related quantization methods like OmniQuant [1], OPTQ [2], AWQ [3], and Mamba-PTQ [4]. Assessing Q-Mamba against these methods would provide better insights into its performance, especially given that Mamba is composed largely of linear layers for which these quantization schemes have proven effectiv

**R3:**In this study, we emphasize that Q-Mamba focuses on the quantization of state caches in Mamba architectures, setting it apart from previous research on quantizing transformers and CNNs. States in Mamba are expanded to be $N$-times larger than standard activations, where $N$ represents the state dimension (128 in Mamba-2 models). Section 4 demonstrates that these state caches contribute significantly to memory consumption, particularly after weights are quantized to low-bit representations. It is important to note that the quantization of linear projections for Mamba models is not the primary contribution of this work, although we present preliminary experiments on Mamba-1 and Mamba-2 using OPTQ[2] and SmoothQuant[5]. On the other hand, Q-Mamba is orthogonal and can be combined with previous quantization methods for linear layers, such as OmniQuant[1], AWQ,[3] and Mamba-PTQ[4].

[1] Shao et al. "OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models", ICLR 2024

[2] Frantar et al. "OPTQ: Accurate Quantization for Generative Pre-trained Transformers", ICLR 2023

[3] Lin et al. "AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration", MLSys 2024

[4] Pierro et al. "Mamba-PTQ: Outlier Channels in Recurrent Large Language Models", ICML 2024

[5] Xiao et al. "SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Model", ICML, 2023

Thank you for your careful reading of our article and providing helpful feedback. We have scrutinized the manuscript and made corresponding modification. We have expanded the introduction to include more details about SSMs and Mamba models ( e.g., the definition of $W_x, W_b$ in Theorem 1) and revised the method section to address and supplement the previously missing information (including Figure 1 and equation for finetuning).

Q1: Decoupled Scale Quantization. According to the SSM equation in (1), the channels are independent from each other. I am not sure why then quantizing the values along the channel dimension would make any difference in the case of multiplication with A & C. Per-state quantization of these matrices and of the states should lead to optimal quantization performance IF these matrix multiplications are in fact done in lower precision. Or the quantization scheme only for compression and matrix multiplications are done in higher precision. There are a lot of quantization details missing here that would explain the merits of the scheme.

R1: During the generation stage, state caches $h_t$ are quantized at each time step to reduce memory consumption and are subsequently dequantized in the next step, which is similar to KV cache quantization. We implement a cuda kernel that load quantized $h_t{-1}$, compute SSM $h_t=A_th_{t-1}+B_tx_t$, then quantize and pack $h_t$ and store into memory. Quantization scale can not be fused into $A_t$ or $C_t$ because $A_t$ is a scalar matrix and $C_t$ only have impact on outputs $y_t$ at the current timestep.

Q2: Efficient Selective Reconstruction. Are the learnt matrices just adjusted to the quantized states? Or they are also quantized using straight-through? What do you learn exactly? Where is the quantization? Do you also learn the quantization scales? Is this a local QAT scheme similar to OmniQuant for transformers? You never introduce an equation of what block-wise reconstruction means.

R2:

ESR aims to adjust parameters for the quantized states, which is a non-trivial task for SSMs. SSMs (and other linear attention mechanisms) rely on the equivalence between two computational modes: the quadratic mode for training and the recurrent mode for inference. This equivalence depends on the linear nature of the hidden-to-hidden transformations, $h_t = A_t h_{t-1} + B_t x_t$. In contrast, RNNs, defined by $h_t = \sigma(A_t h_{t-1} + B_t x_t)$ (where $\sigma$ is a nonlinear activation function) cannot be parallelized ( known as backpropagation through time (BPTT)) . To minimize memory bandwidth utilization, we store state caches as low-bit elements and subsequently load and dequantize them before computation at the next timestep. This process introduces a new sequence transformation, defined as $h_t^q = \bar{A} Q(h_{t-1}^q) + \bar{B} x_t$, where $Q(\cdot)$ is a non-linear quantization function. It is crucial to distinguish $h_t^q = \bar{A} Q(h_{t-1}^q) + \bar{B} x_t$ from the quantized value of the original $h_t$, i.e., $Q(h_t) = Q(\bar{A} h_{t-1} + \bar{B} x_t)$. A naive approach would directly apply recurrent mode for token-by-token generation in the training. However, this method becomes computationally prohibitive for large input lengths (e.g., 2048).

To address this, we propose Efficient Selectivity Reconstruction (ESR), which utilizes the parallel algorithm to compute $h_t$ across all timesteps and simulates quantization errors by quantizing only a single step during training. We employ straight-through estimation for the backpropagation of the quantization function $Q$. In this process, we freeze the linear projections corresponding to the SSM inputs $x$ and $z$, while keeping the linear projections for selective parameters $B$, $C$, and $\Delta$ learnable (account for only about 2% of the total parameters). So far, we do not make the quantization scales learnable. Specifically, we have updated Figure 1 to highlight the learnable parameters and added a new equation to define this optimization process in Section 5.2.2

Q3: Figure 2: is this the memory consumption during pre-fill or generation? Figure 6, the legend hides the graph. Figure 9, the caption does not explain what the different colours are.

R3: Figure 2 shows the memory consumption during generation. Because of time limitation, we will modify Figure 6 and Figure 9 in the final papers.

Q1: The background section could be enhanced for clarity for readers who are not familiar with Mamba. It would be beneficial to include additional information about the differences between the Mamba-1 and Mamba-2 models.

R1: Thank you for your helpful feedback. We have sigificantly modified this section in the revised paper.

Q2: Some ablation studies on hyperparameters, such as optimizer settings ( learning rate, and weight decay) should be provided to demonstrate the robustness of ESR.

R2: For all size models, we set the same hyperparameters. We utilize the AdamW optimizer with zero weight decay to optimize the learnable parameters in ESR. The learning rate for the learnable parameters is set to 1E-3. In Table 1, we demonstrate that different learning rates can all lead to improvements in performance.

Table 1 Impact of learning rate.

Q3: Considering the need for fast deployment, the total time consumption for layer-wise construction should be included, particularly for the largest Mamba models.

R3: ESR only needs a single A800 80G GPU and with only 128 training samples. The training time takes about 8 hours for the Mamba2-2.7B.

Q4: The appendix provides the results of the GPTQ[1] and SmoothQuant[2] methods on Mamba models. Could you provide more results for different quantization settings (e.g., W3A16, W4A4)? I am curious about the bit precision that can be achieved with Mamba model quantization.

R4: Table 2 presents experimental results on various weight bit configurations for the Mamba2-370M model using the GPTQ [1] method. Compared to Transformers, Mamba models exhibit more significant performance degradation in extremly low-bit quantization. This issue may be attributed to large outliers in the inputs to the output projection layer, as shown in Figure 7.

Table 2 Experiments on various weight bits for Mamba2-370M with GPTQ[1] method.

[1] Frantar et al. "OPTQ: Accurate Quantization for Generative Pre-trained Transformers", ICLR 2023

We sincerely appreciate your reminder and have now submitted the revised version of the paper. We understand your concern regarding the timing, and we apologize for any inconvenience this may have caused. According to the official notification, the discussion period has been extended to December 2nd (a six-day extension). Thank you again for your valuable feedback. If you have any additional questions or suggestions, please let us know

As of this time, the submission does not have a revision attached. In addition, responding to our reviews on the last day of the discussion phase does not give us time to assess your responses properly.