Q -Mamba: Towards more efficient Mamba models via post-training quantization (2025)

Abstract: State Space Models (SSMs), such as Mamba, have recently demonstrated potential in language understanding tasks, positioning them as competitors to transformer architectures. However, our investigations reveal that the Mamba architecture still has room for further optimization—not only in linear projections but also in state caches, which contribute significantly to memory consumption, particularly after quantizing the former into low bits. After a theoretical analysis of the causes of outliers in states, we propose Decoupled Scale Quantization (DSQ), which mitigates outliers in both the state and channel dimensions by applying separate quantization scales. To preserve the selective ability of quantized Mamba, we introduce Efficient Selectivity Reconstruction (ESR), a novel quantization simulation scheme in block-wise reconstruction that enables fast parallel scan algorithms with the non-linear quantization function. We demonstrate the effectiveness of Q-Mamba across various quantization settings, model sizes, and both generation and zero-shot tasks. In particular, for Mamba2-2.7B with W8A8H4 (8-bit weights and activations, 4-bit state caches) quantization, Q-Mamba achieves a 50% reduction in memory consumption with only a 2.13% average accuracy degradation on zero-shot tasks.

## introduction large language models (llms), such as llama @xcite and @xcite , have shown exceptional capabilities in generalpurpose language understanding @xcite @xcite . however, llms based on transformer architectures still face a significant limitation: the compu-tational cost of their attention mechanism scales quadratically with the sequence length @xcite . therefore, prior works have focused on more efficient attention variants, such as structured state space models (ssms) @xcite @xcite and linear attention @xcite @xcite . among these, the mamba architecture @xcite has been shown to match or exceed the downstream accuracy of transformers on standard language modeling tasks @xcite . following its success in natural language understanding, it has also garnered significant attention in other research areas, such as vision and multimodal tasks @xcite @xcite . like transformers, mamba language models also operate in two computation phases @xcite . the first is the prefill phase, where all input prompt tokens are processed in parallel through the model's forward pass to generate the first output token. during this phase, mamba models @xcite ) employ a hardware-efficient parallel algorithm to compute ssms (section 3). the second is the token generation phase, where subsequent output tokens are generated sequentially, relying on the cached state from previous tokens in the sequence. due to the lack of computational parallelism, this phase tends to be more memory-bound and contributes significantly to the total generation latency. although mamba has successfully replaced the o(t 2 ) attention module with o(t ) selective state space models, our profiling results in section 4 indicate that it still suffers from two inefficiencies during the generation stage. firstly, similar to transformers, the mamba architecture consists of large linear layers, which require substantial gpu memory and slow down token generation (figure 2b ). secondly, as larger states allow more information to be stored, states in mamba are expanded to be n times larger than vanilla activations, where n is the state dimension (128 in mamba-2 models). consequently, these state caches account for a significant portion of memory consumption, especially after quantizing weights to low bits (79.6% in mamba2-2.7b with a batch size of 128, as shown figure 1 : schematic of the ptq framework for mamba. left: the selective parameters b, ∆, and c, along with the ssm inputs x, are generated by the input projections in the mamba block. middle: after quantizing states using dsq, esr updates a small number of selective parameters (approximately 2% of the total) in a block-wise reconstruction manner. right: finally, we quantize the linear projection into w8a8. in figure 2a ). in this paper, we address a key question: can mamba models be further optimized through model compression techniques? in this paper, we propose q-mamba, which quantizes both linear projections and state caches into low-bit integers for mamba models. although previous research has successfully quantized key and value (kv) caches into low-bit representations in transformers @xcite @xcite , this work is the first to explore the quantization of state cache in mamba architectures. we observe that states exhibit both outlier channels and outlier states (i.e., the state dimension contains large values across all channel dimensions), as shown in figure 3 . further theoretical analysis reveals this phenomenon results from the computation of the outer products of two activations, each contains outliers in distinct dimensions. this observation motivates us to propose decoupled scale quantization (dsq), which utilizes separate quantization scales for both dimensions. additionally, the non-linear nature of the quantization function disrupts the original equivalence between recurrence and quadratic dual form, the latter being essential for efficient training. to address this, we propose efficient selectivity reconstruction (esr), which simulates quantization errors by quantizing only the final timestep during training. specifically, esr updates a small number of selective parameters (approximately 2% of the total) using just 128 training samples in a block-wise reconstruction manner. extensive experiments demonstrate that our methods achieve significant performance improvements for mamba families on various evaluation metrics. to the best of our knowledge, we are the first to achieve w8a8h4 (8-bit linear projection and 4-bit states) for the mamba architectures. for generation tasks, q-mamba achieved perplexities of 12.99 and 16.90 with 4-bit states on wikitext2 @xcite and c4 @xcite , respectively, while baseline methods degraded to 21.18 and 29.86 even with 6-bit quantization. additionally, q-mamba achieves w8a8h4 quantization for zero-shot tasks with only 2.13% and 2.11% average accuracy degradation on mamba2-2.7b and mamba2-1.3b, respectively. ## foundations state space model. state space models (ssms) are a recent class of sequence models for deep learning inspired by a particular continuous system in equation (2). it maps a 1-dimensional input sequence x @xmath0,1) : ## conclusion in this paper, we propose q-mamba, a novel quantization framework designed for state caches in mamba models. we conduct a theoretical analysis of outliers in states and propose decoupled scale quantization (dsq) which decouple scales in the state and channel dimensions, dsq mitigates outliers in both dimensions while introducing negligible overhead. to further boost performance through block-wise reconstruction, we propose efficient selectivity reconstruction (esr), which includes a novel quantization simulation method that enables efficient fine-tuning of selective parameters with parallel scan mode. in conclusion, q-mamba demonstrates that mamba architectures have the potential for further optimization when combined with other model compression techniques. limitations in this paper, we propose q-mamba, a novel quantization framework designed for state caches in mamba models. although q-mamba reduces memory usage by half and achieves a 1.18x speedup on gpus with only a 2.13% average accuracy degradation on zero-shot tasks, maximizing acceleration in large language models llms based on mamba via quantization requires greater hardware support. previous research on hardware accelerators has primarily focused on llms (wang et al., 2021; sridharan et al., 2023) based on transformer architectures. we hope this paper will inspire more researchers to focus on developing customized hardware for low-bit ssm models. rowan zellers, ari holtzman, yonatan bisk, ali farhadi, and yejin choi. 2019. hellaswag: can a machine really finish your sentence? in proceedings of the 57th conference of the association for computational linguistics, acl 2019, florence, italy, july 28-august 2, 2019, volume 1: long papers, pages 4791-4800. association for computational linguistics. lianghui zhu, bencheng liao, qian zhang, xinlong wang, wenyu liu, and xinggang wang. 2024. vision mamba: efficient visual representation learning with bidirectional state space model. corr, abs/2401.09417.