Bi-Mamba: Towards Accurate 1-Bit State Space Models

典型的Mamba选择性状态空间模型（SSM）解决了Transformer的几个限制，例如序列长度的二次计算复杂度以及由于键值缓存导致的显著推理时间内存需求。然而，Mamba模型不断增长的规模继续带来训练和部署的挑战，并因巨大的能源消耗引发了环境问题。在这项工作中，我们引入了Bi-Mamba，这是一种可扩展且强大的1位Mamba架构，旨在更高效地构建多个规模的大型语言模型，包括7.8亿、13亿和27亿参数。Bi-Mamba模型从头开始在与常规大型语言模型相同的数据量上进行训练，使用自回归蒸馏损失。广泛的实验结果表明，Bi-Mamba在语言建模任务上的表现与其全精度对应模型（如FP16或BF16）相当，且比后训练二值化（PTB）的Mamba基线模型具有更高的准确性，同时显著减少了内存占用和能源消耗。我们的研究开创了一种低比特表示下的线性计算复杂度大型语言模型框架，并促进了未来为高效1位Mamba基大型语言模型设计专用硬件的发展。

该论文旨在解决大型语言模型（LLM）在训练和部署过程中面临的计算复杂度高、内存需求大以及能源消耗大的问题。这些问题随着模型规模的不断增大而变得尤为突出。

论文提出了一种名为Bi-Mamba的1位Mamba架构，通过将模型参数量化为1位来显著减少内存占用和能源消耗，同时保持与全精度模型相当的性能。这种低比特表示方法不仅解决了Transformer模型的二次计算复杂度问题，还为未来专门硬件的设计提供了新的方向。

论文通过广泛的实验验证了Bi-Mamba模型在语言建模任务上的有效性，展示了其在不同规模下的性能与全精度模型相当，并且优于后训练二值化（PTB）的Mamba基线模型。此外，该研究还强调了Bi-Mamba模型在实际应用中的环境友好性，为高效能、低能耗的LLM设计提供了新的思路。论文中提到的模型大小包括780M、1.3B和2.7B参数量的版本。目前没有提及是否有开源代码。

近年来，关于提高LLM效率的研究层出不穷。例如，'Sparse Transformers' (Child et al., 2019) 提出了稀疏注意力机制以降低计算复杂度；'EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks' (Tan & Le, 2019) 探讨了模型缩放的有效策略；'Leveraging Sparsity in Transformer Models' (Gordon et al., 2020) 则研究了如何利用稀疏性来优化Transformer模型。这些研究与Bi-Mamba的目标相似，都是为了提高模型的效率和可扩展性。