Bi-Mamba: 迈向精确的 1 比特状态空间模型

Mamba 中使用的典型选择性状态空间模型 (Selective State-Space Model, SSM) 解决了 Transformer 的若干局限性，例如相对于序列长度的二次计算复杂度，以及由于键值 (key-value, KV) 缓存导致推理期间显著的内存需求。然而，Mamba 模型规模的不断增长持续给训练和部署带来挑战，这主要是由于它们在训练和推理期间都有大量的计算需求。在本工作中，我们引入了 $\texttt{Bi-Mamba}$，这是一种可扩展且强大的 1 比特 Mamba 架构，旨在实现更高效的大型语言模型 (LLM)，模型规模分别为 780M、1.3B 和 2.7B 参数。$\texttt{Bi-Mamba}$ 模型在标准 LLM 规模数据集上从头开始训练，使用自回归蒸馏损失。在语言建模基准上的大量实验表明，$\texttt{Bi-Mamba}$ 实现了与其全精度 (FP16 或 BF16) 对应模型相当的性能，同时优于训练后二值化 (post-training binarization, PTB) Mamba 和二值化感知训练 (binarization-aware training, BAT) Transformer 基线。此外，与原始 Mamba 相比，$\texttt{Bi-Mamba}$ 大幅降低了内存使用和计算成本。我们的工作开创了低比特表示下线性复杂度 LLM 的新方向，并为针对高效的基于 1 比特 Mamba 模型优化的专用硬件设计铺平了道路。代码和预训练权重可在 https://github.com/Tangshengku/Bi-Mamba 获取。