点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
典型的Mamba选择状态空间模型（SSM）解决了 Transformer 的一些限制，例如序列长度和关键值缓存的二次计算复杂度，以及由于推理时间中的大量内存需求。
然而，Mamba模型的不断增长尺寸仍然带来了训练和部署挑战，并由于大量的能源消耗引发了环境问题。
在这项工作中，作者提出了双Mamba，这是一种可扩展和强大的1比特Mamba架构，旨在为具有多种尺寸（780M、1.3亿和2.7亿）的更高效的大型语言模型。双Mamba模型从零开始在数据量上进行训练，使用自回归蒸馏损失。
在语言建模方面的广泛实验结果表明，双Mamba在性能上与全精度对应物（例如FP16或BF16）相当，并且比Mamba后训练二进制化（PTB） Baseline 实现更高的准确性，同时显著降低了与原始Mamba模型相比的内存占用和能源消耗。
作者的研究开创了在低比特表示下的线性计算复杂度大型语言模型框架，并有助于未来设计专为高效1比特Mamba基础的大型语言模型而设计的专用硬件。
unset
unset1 Introductionunset
unset
选择性状态空间模型（SSM）最近在语言建模中作为一种强大的替代方案，成为Transformer（Vaswani等人，2017年）的有力竞争者，在小型到中等规模上表现出可比或更好的性能。SSM在训练过程中具有线性缩放的时间序列长度，而在生成过程中具有恒定的状态大小，这显著降低了计算和内存开销，使其在速度和内存使用方面更加高效。
Mamba 的代表性自注意力模型在处理长序列时具有显著优势，因为其线性复杂度。相比之下，传统 Transformer 在序列长度增加时具有二次复杂度。这意味着在涉及大型输入序列或扩展上下文的情况下，Mamba更高效，因为其内存和计算需求与序列长度成线性比例。在实际应用中，这使得Mamba能够更快地处理长序列，并降低资源消耗，使其成为处理长文档、对话 Agent 以及任何需要管理大量上下文信息场景的理想选择。另一方面， Transformer 在上下文长度增加时需要指数级更多的资源，这可能导致长上下文任务中的 Bottleneck 。
近年来，Transformers的早期研究已经得到了广泛的关注。为了解决其高计算和存储成本问题，提出了许多方法，如剪枝，模型量化，以及KV Cache 。在这些方法中，量化被证明是非常有效的（Dettmers等人，2022年）。研究行人成功地将从16位到8位、4位和甚至1位的Transformer模型进行了量化，通常性能降级最小。
然而，关于SSM模型或Mamba在低位量化或甚至二进制化后的行为的研究还很少。
在本文中，作者引入了Bi-Mamba，一种通过将权重减少到二进制设置来应用极端量化的方法。通过这种极端量化，作者证明了SSM模型可以在训练和推理过程中有效地二进制化，同时保持高性能，同时显著降低内存占用和能耗。
作者的工作首次提出了一种在线性计算复杂性模型中使用1比特表示的新框架，这可能有助于开发专门针对全低比特基于短语表示模型（SSM）的专用硬件。作者提供了全面的实验结果，表明Bi-Mamba在与其他全精度对应物LLM的比较中实现了竞争性能，从而证明了将SSM模型二进制化的可行性和优势。
为了进一步了解二进制化对Mambas行为的影响，作者对Mambas中预训练和二进制化权重的分布进行了广泛的实证研究。本研究的成果详细阐述了实验和附录部分，得出了两个关键观察：
如图2所示，在训练后二进制化方法如BiLLM（黄等人，2024年）和PB-LLM（商等人，2023年）通常会导致在二进制化后Mamba上的权重分布发生偏移，从而导致与最优二进制权重分布不匹配。
然而，作者的二进制化感知的训练确保了二进制化权重仍然接近原始权重分布，保留了二进制表示的最大能力。
根据作者的实证实验，应用现有的LLM后binarization方法（Frantar等，2023；黄等，2024），即使保留显著权重，通常也会严重降低Mamba模型的性能。如果不考虑显著权重，后binarization感知训练似乎是唯一可行的解决方案，既能有效地binarize像Mamba这样的模型，又能保持竞争性能。
在本文中，作者的贡献如下：
(1) 这是首次成功将Mamba架构二进制化至1比特，同时保持强大性能的工作。
(2) 作者探索了Mamba二进制化的潜在参数空间，为未来的研究提供了宝贵见解。
(3) 由作者方法预训练的Bi-Mamba模型可以在资源受限的场景中作为强大的基础模型，也可轻松适应其他NLP应用。
unset
unset2 Related Workunset
unset
PTQ 。由于大型语言模型（LLMs）参数数量庞大，运行起来资源消耗巨大。因此，在推理过程中，它们通常会被量化到低位表示。这种量化通常称为PTQ（PTQ），其中量化操作应用于预训练模型。LLMs的PTQ方法通常基于以下经验观察：预训练Transformer中的少数显著特征具有极大的幅度，这些特征的量化需要特别小心。PTQ方法导致了INT8权重和激活量量的近似无损失准确率下降，以及INT4权重仅量化的方法。尽管许多工作专注于将PTQ权重量化到4位以下，但通常会导致模型性能急剧下降。
量化感知训练 。与后量化训练不同，量化感知训练（QAT）旨在在训练过程中学习量化模型。在LLM时代，与后量化训练相比，量化感知训练（QAT）的研究较少，因为LLM通常需要大量的计算资源进行训练。李等人（2023年）提出了一种在量化感知训练中采用LoRA微调的方法，使得量化过程具有资源效益。然而，这些方法仍然需要使用预训练权重来初始化学生LLM。沃茨曼等人（2023年）提出了一种从随机初始化中稳定低比特大视觉语言模型训练的方法。最近，马等人（2024年）使用十进制权重从随机初始化中训练LLM，实现了非平凡的性能。马等人（2024年）进一步提出了一种基于蒸馏的方法，从零开始预训练二进制LLM。
虽然语言模型通常使用Transformer（Vaswani等人，2017年）构建，但它们的自注意操作存在二次时间复杂度问题，这使得Transformer在处理长序列时效率低下。因此，已经进行了大量的研究，旨在解决这个问题。在这些努力中，状态空间模型（SSMs）是一种循环神经网络。最新的高级SSM，例如Mamba（Gu和Dao，2024年）和Mamba-2（Dao和Gu，2024年）与Transformer相比，在序列长度方面表现出优越的性能，同时具有线性复杂度。尽管它们具有有前景的能力，但如何有效地量化这种架构的研究很少。在作者工作的同时，Pierro和Abreu（2024年）研究了Mamba的post-training量化。
unset
unset3 Approachunset
unset
Preliminary: Mamba Series
Mamba属于一类被称为状态空间模型（SSMs）的模型，它能够在保持实用性的同时，提供与Transformer相当的表现和扩展定律，前提是序列长度非常长（例如，一百万个 Token ）。它通过消除注意力机制中的“二次 Bottleneck ”来实现这一点。常规的结构状态空间序列模型（S4）将一个一维序列 转换为另一个序列 ，使用潜在状态 如下所示：
, , , 。 是隐藏状态， 是输入，模型每次获得的观测值。 然后表示隐藏状态的导数，即状态如何变化。
由于方程1中的参数在时间上是恒定的，因此该模型是线性时不变（LTI）的。LTI模型在处理序列时，对所有元素给予相同的关注，这使得模型无法有效地理解自然语言。为了解决这个问题，Mamba对其进行了改进，将其转变为选择状态空间模型，其中B和C是通过x_{t}的线性投影得到的。同时，上述方程适用于具有连续输入和输出信号的动态系统。因此，在处理文本序列时，需要进行离散化。
然后，选择性SSM的计算过程如下： 与 和一个可学习的参数相关。
围绕选择性SSM，结合线性投影进行输入和输出以及 LayerNorm ，形成了Mamba的基本模块。
与Mamba相比，Mamba-2进一步用状态空间对偶（SSD）替代了选择性SSM。在SSD中，A被简化为标量乘以单位矩阵的结构，而在Mamba中，SSD使用了一个较大的头维数，该维数在Mamba中等于1。同时，Mamba-2块也引入了一些简化措施，例如在Mamba块中移除了顺序线性投影，以促进并行训练。Mamba-2的基本结构如图2所示。
Bi-Mamba
首先，作者需要确定Mamba架构中哪些权重矩阵可以进行二进制化。作者以最新的Mamba-2（Dao和Gu，2024年）作为基础架构。根据上述描述，作者可以更直观地理解，，，和的SSD矩阵，以及Mamba-2中的其他层，包括_embedding，layer normalization (LN)，Conv-1d，内层线性投影，外层线性投影_，以更好地了解二进制化后的效果，从而确定二进制化空间。简而言之，是转换状态矩阵，表示当前状态如何过渡到下一个状态。它回答了模型应该如何在时间上逐渐忘记状态的较少相关部分？
将新输入映射到状态，解决了模型应该保留新输入的哪些部分？的问题。
将状态映射到SSM的输出，询问模型如何利用当前状态进行准确的下一个预测？。
显示新输入如何直接影响输出，作为修改的 Shortcut ，询问模型如何将新输入纳入预测？。
考虑到作者的基础架构Mamba-2，作者在图2和表1.2中展示了不同大小下不同层的具体参数比例。在Mamba-2中，大多数模型参数都在线性模块中，排除因果头。例如，在Mamba-2-2.7B中，这些参数占总模型的95.2%。此外，嵌入模块与因果头共享参数。对嵌入进行二进制化会显著降低其表示 Token 语义的能力，从而降低模型性能。
因此，在作者的双Mamba中，作者只二进制化线性模块（不包括因果头）的参数。这种策略旨在保持90%的高压缩比，同时仍然允许二进制化模型有效运行。
简化的可学习缩放因子以提高更好的性能 为了二进制化Mamba，作者用FBI-LLM（Ma等人，2024a）提出的FBI-Linear模块替换原始的线性模块。具体而言，FBI-Linear模块主要由两部分组成：仅包含{1,-1}的矩阵和高精度的缩放因子和。FBI-Linear的推理过程如下：
其中， 的计算方法是：先对 进行列向量乘法，然后与 相加：
在第i层中，和是可学习的缩放和偏移因子。
训练目标。作者的训练目标是目标学生模型输出与预训练教师模型在每个自回归方案中的输出的交叉熵损失，用于下一个 Token 预测。这可以表示为：
在这里， 表示输入 Token 的数量。在这里， 表示由教师模型在步骤 上预测到的词汇表上的 Token 分布，而 是学生模型相应地预测到的分布。
整体设计双Mamba。在自回归蒸馏（Ma等人，2024年）的binarization-aware训练中，作者计算高精度预训练模型和目标双Mamba模型输出概率分布之间的交叉熵。在此过程中，和是可学习的参数，而使用可学习的高精度矩阵中的函数导出。由于函数不可导，作者使用直线估计器（STE）(Bengio等人，2013年)作为先前的研究（Rastegari等人，2016年；Alizadeh等人，2019年）在BNNs中近似输入变量的梯度，从而使得反向传播得以继续进行。
unset
unset4 Experimentsunset
unset
在作者的实验中，作者仅在最具线性的模块中二进制化参数，同时保持其他模型参数和激活的原精度。值得注意的是，作者并没有严格地将二进制化参数表示为1-bit，而是使用高精度值来模拟二进制化参数。作者在不同的尺度上训练Bi-Mamba，并在多个任务上评估其性能。
Setup
数据集 遵循FBI-LLM的方法，作者使用Amber数据集（Liu等人，2023年）进行训练。该数据集总共包含1.26万亿个 Token （tokens）来自RefinedWeb（Penedo等人，2023年）、StarCoder（Li等人，2023年）和RedPajama-v1（计算机，2023年）。数据被分割成360个块，每个块平均包含约35亿个 Token 。
训练细节：作者在Mamba-2架构的不同尺度上训练Bi-Mamba。具体来说，作者在780M、1.3B和2.7B Mamba-2模型中将输入投影和输出投影矩阵二进制化。作者将LLaMA2-7B作为所有Bi-Mamba的教师模型，用于计算自回归蒸馏损失。因此，作者训练的所有Bi-Mamba具有与LLaMA2相同的词汇和分词器。作者总共使用32块NVIDIA A100 GPU进行训练，并在训练过程中保持BF16精度。
评估指标
作者根据模型在包括BoolQ（Clark等人，2019年）、PIQA（Bisk等人，2020年）、HellaSwag（Zellers等人，2019年）、Winogrande（Sakaguchi等人，2021年）、ARC（Clark等人，2018年）和OpenbookQA（Mihaylov等人，2018年）等下游任务中的零样本表现来评估模型。所有下游评估都使用_lm-evaluation-hamess_包（Gao等人，2024年）。作者还使用Wikitext2（Merity等人，2016年）、PTB（Marcus等人，1993年）和C4（Raffel等人，2020年）数据集作为评估指标。困惑度衡量一个概率模型对一个 Token 的预测能力，从而定量衡量模型的生成能力。
Baseline 方法 作者将作者的工作与量化（quantization）和二进制化（binarization）方法进行比较，这些方法包括 GPTQ（Frantar 等，2023）和 Bi-LLM（Huang 等，2024）。GPTQ 是一种PTQ方法，而 Bi-LLM 是一种后训练二进制化方法。作者使用 GPTQ 的官方实现，并将 Mamba-2 模型分别量化为 3 和 2 位。对于 Bi-LLM，作者也使用他们的官方实现，并将 Mamba-2 模型二进制化。此外，作者还包括各种大小的开源全精度 Transformer 模型作为参考，如 OPT（Zhang 等，2022）和 TinyLLaMA（Zhang 等，2024）。
Main Results
表2展示了作者的Bi-Mamba模型在下游任务和Wiki2，PTB和C4数据集上的性能对比各种 Baseline 。这些评估有助于了解模型在不进行进一步的任务特定微调时的泛化能力。性能比较的图表展示如图3所示。
对于780M的Mamba-2模型，Bi-Mamba的平均下游性能达到44.5，优于GPTQ-3bit和Bi-LLM的42.6和38.1，分别。在困惑度评估中，Bi-Mamba在Wiki2、PTB和C4上的得分分别为14.2、34.4和15.0，而 Baseline 模型在困惑度上可达10倍以上。对于1.3B模型，Bi-Mamba的平均下游准确性达到45.8，超过了GPTQ-2bit的35.3和Bi-LLM的36.2。这一性能表明，在模型大小更广泛的范围内，Bi-Mamba的泛化能力得到提升。此外，Bi-Mamba的困惑度也有显著改善，在Wiki2、PTB和C4数据集上的得分分别为13.2、30.8和14.0，而GPTQ-2bit和Bi-LLM的困惑度明显更高，说明Bi-Mamba的二进制化在保持语言连贯性方面更具效率。
对于2.7B模型，Bi-Mamba进一步扩大了领先优势，平均下游准确性达到49.3，而GPTQ-2bit和Bi-LLM分别为35.8和38.5。值得注意的是，Bi-Mamba在所有数据集上保持了低困惑度，Wiki2、PTB和C4数据集上的得分分别为10.7、24.2和11.9，这些结果突显了Bi-Mamba在模型复杂度提高时，既保持了任务准确性又保持了语言流畅性。
总之，Bi-Mamba在模型大小和大小上均展现出优越的零样本性能和困惑度降低，证实了其相对于GPTQ和BiLLM的稳健性和多样性，尤其是在大型、更复杂的模型中。
unset
unset5 Analysisunset
unset
Training Result Dynamics
在本节中，作者讨论了随着不同训练成本/预算的训练进程，双Mamba的性能。主要结果如图4所示。作者在Mamba-2-780M模型上提供了不同训练成本下的下游性能和困惑度曲线。2.7B和1.3B Mamba-2模型上的更多结果请参见附录。从图中，作者可以观察到训练开始时困惑度下降较快，逐渐收敛到全精度困惑度。双Mamba在维基2和C4数据集上的困惑度比PTB数据集上的困惑度更稳定。值得注意的是，在C4数据集上，双Mamba的最终困惑度甚至低于全精度模型，突显了双Mamba的优越性能。有趣的是，在训练初期，作者的双Mamba模型在困惑度上超过了GPTQ-3bit，证明了二进制感知训练的有效性。由于GPTQ-2bit和BiLLM在所有数据集上的困惑度极高，作者在图中忽略它们，并在表2中找到详细的结果。
转向下游任务评估，作者首先观察到二进制模型的灾难性性能下降，其性能在许多基准测试（如ARC-E和ARC-C上的结果）甚至低于随机结果。这表明直接应用简单的二进制化方法破坏了全精度模型的能力。然而，经过二进制感知训练后，模型在所有基准测试上的性能得到恢复，并最终超过了包括GPTQ-3bit、GPT-2bit和Bi-LLM在内的所有 Baseline 。这强调了在实现竞争力的结果方面，二进制感知训练的重要性。
Binarization Space
在本节中，作者探索了Mamba模型的二进制化空间，并讨论了二进制化每个部分的影响。作者进行了将In_Proj和Out_Proj单独二进制化和同时二进制化的实验。二进制化后的模型使用相同的数据进行训练。结果如表3所示。部分和全二进制化与全精度Mamba模型在困惑度及下游任务上的性能进行了比较。
首先，表中的结果表明，部分二进制化通常保留的全零样本准确率高于全二进制化，尽管它仍然显示出与原始Mamba-2模型相比的性能下降。然而，部分和全二进制化模型之间的性能差距并不显著。例如，在Mamba-2-2.7B模型中，部分二进制化平均达到53.1的准确率，而全二进制化将此降低到49.3。
在所有模型大小中，部分二进制化的Bi-Mamba在大多数基准测试上始终优于全二进制化的Bi-Mamba，这表明部分二进制化可以减轻全二进制化引起的某些准确性损失。这也暗示着，虽然部分二进制化保留了一定的优势，但全二进制化仍然具有高度竞争力，并未显著落后。
在困惑度方面，全二进制化模型与部分模型表现相当。例如，在Mamba-2-780M模型中，全二进制化Bi-Mamba（Fully）的C4数据集困惑度为15.0，而部分二进制化Bi-Mamba（Partially）为14.4，表明全二进制化并未显著增加困惑度。这些发现表明，全二进制化模型可以在部分二进制化模型保持竞争力的同时，仍可充分利用更大的存储和计算效率。结果表明，全二进制化，与最初预期相反，并未严重影响模型有效性，提供了一种可行、高效的选择。
Storage Efficiency
模型二进制化可以显著减少磁盘存储需求。遵循Bi-LLM (Huang等人，2024年)，作者提供了与全精度模型相比，在不同的模型大小下，作者的Bi-Mamba的理论存储需求，如表4所示。对于每个参数大小，原始全精度Mamba-2模型的存储需求明显大于包括部分和全二进制化的binarized Bi-Mamba的存储需求。具体而言，完全二进制化的Bi-Mamba展示了最高的压缩比，实现了80%以上的减少。相比之下，部分二进制化的Bi-Mamba提供了相对适中的压缩，范围在55%到60%之间。这种分析突显了完全二进制化在显著降低存储需求的同时保持模型参数数量，使其成为大型模型的高效存储选择。
unset
unset6 Conclusionunset
unset
作者提出了Bi-Mamba，一种可扩展且高效的1-bit Mamba架构，专为多种大小的语言模型设计：780M、1.3亿和2.7亿参数。
作者首先在Mamba架构内确定二进制化空间。
然后，从零开始在大型数据集上训练Bi-Mamba模型，类似于标准的LLM预训练，使用自回归蒸馏损失。大量语言建模实验表明，Bi-Mamba在性能上略低于其全精度对应物（例如，FP16或BF16），但相比原始精度Mamba，显著降低了内存使用和计算成本。
本研究提供了一个新颖、可访问的第一种低位框架，线性计算复杂度，为开发针对高效1-bit Mamba架构的专用硬件奠定了基础。
[0]. Bi-Mamba: Towards Accurate 1-Bit State Space Models.
点击上方卡片，关注 「AI视界引擎」 公众号