Bi-Mamba用1比特压缩模型,实现高效大语言模型部署。
Bi-Mamba: Towards Accurate 1-Bit State Space Models
- 采用自回归蒸馏训练1比特Mamba,避免传统二值化损失。
- 780M~2.7B参数模型性能接近全精度,且内存与计算大幅降低。
- 适合追求低资源推理的部署场景,推动专用硬件设计。
典型的选择性状态空间模型(SSM)如Mamba,克服了Transformer在序列长度上二次复杂度和推理时键值缓存占用大量内存的问题。然而,随着Mamba模型规模增大,训练与部署仍面临巨大计算压力。本文提出可扩展、高性能的1比特Mamba架构Bi-Mamba,支持780M、1.3B和2.7B参数规模。模型在标准大语言模型数据集上从头训练,使用自回归蒸馏损失。在语言建模基准上的实验表明,Bi-Mamba性能接近同规模全精度(FP16/BF16)模型,优于后训练二值化(PTB)Mamba和二值化感知训练(BAT)Transformer基线。同时,相较于原始Mamba,Bi-Mamba显著降低内存占用与计算开销。本工作开创了低比特表示下线性复杂度大语言模型的新方向,并为面向1比特Mamba的专用硬件设计提供路径。代码与预训练权重已公开于https://github.com/Tangshengku/Bi-Mamba。
原文摘要 · Abstract (English)
The typical Selective State-Space Model (SSM) used in Mamba addresses several limitations of Transformers, such as the quadratic computational complexity with respect to sequence length and the significant memory requirements during inference due to the key-value (KV) cache. However, the increasing size of Mamba models continues to pose challenges for training and deployment, particularly due to their substantial computational demands during both training and inference. In this work, we introduce $\texttt{Bi-Mamba}$, a scalable and powerful 1-bit Mamba architecture designed to enable more efficient large language models (LLMs), with model sizes of 780M, 1.3B, and 2.7B parameters. $\texttt{Bi-Mamba}$ models are trained from scratch on a standard LLM-scale dataset using an autoregressive distillation loss. Extensive experiments on language modeling benchmarks demonstrate that $\texttt{Bi-Mamba}$ achieves performance comparable to its full-precision (FP16 or BF16) counterparts, while outperforming post-training binarization (PTB) Mamba and binarization-aware training (BAT) Transformer baselines. Moreover, $\texttt{Bi-Mamba}$ drastically reduces memory usage and computational cost compared to the original Mamba. Our work pioneers a new line of linear-complexity LLMs under low-bit representation and provides the way for the design of specialized hardware optimized for efficient 1-bit Mamba-based models. Code and the pre-trained weights are available at https://github.com/Tangshengku/Bi-Mamba.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。