Mamba State-Space Models Can Be Strong Downstream Learners

最近，Mamba状态空间模型（SSM）在各种任务中表现优于最先进的Transformer大语言模型（LLM）并被广泛采用。然而，Mamba的下游学习能力仍未被探索（例如，混合精度（MPFT）和参数高效微调（PEFT）），或者被低估（例如，在上下文学习（ICL）方面）。关于后者，最近的研究报告称，Mamba的ICL与非标准基准下的SOTA Transformer LLM相媲美。相比之下，我们展示了在标准基准下，预训练的Mamba模型仅实现了与可比较的Transformer相比38％的ICL性能提升（超过零-shot）。 在Mamba架构中启用MPFT和PEFT具有挑战性，因为涉及到循环动态和高度定制的CUDA内核。然而，我们证明了Mamba的循环动态对小的输入变化是强健的，使用动态系统理论进行了证明。实验证明，由于混合精度，Mamba推理和微调的性能变化与Transformer LLM相一致。此外，我们展示了针对Mamba的关键内存缓冲区进行低秩适应的定制CUDA内核，可以规范SSM参数，从而实现参数效率并保持加速。我们展示了将MPFT和PEFT结合起来，可以实现每秒最多2.15个标记和每个标记的内存减少65.5％，同时实现与相应微调的Transformer相比81.5％的ICL性能提升（超过零-shot）。

本论文旨在探索Mamba状态空间模型（SSM）在深度学习中的下游学习能力，包括混合精度（MPFT）和参数高效微调（PEFT）以及上下文学习（ICL），并与目前最先进的Transformer大型语言模型（LLMs）进行比较。

本论文提出了一种基于动力系统理论的方法来证明Mamba的循环动力学对小输入变化具有鲁棒性，并且在Mamba的自定义CUDA内核中针对关键内存缓冲区进行低秩适应，从而实现参数效率和保留加速的同时，使得混合精度和参数高效微调相结合。

本论文的亮点包括：1.提出了一种基于动力系统理论的方法来证明Mamba的循环动力学对小输入变化具有鲁棒性；2.通过针对关键内存缓冲区进行低秩适应，实现了参数效率和保留加速；3.结合混合精度和参数高效微调，实现了更高的tokens-per-second和更低的per-token-memory，同时保留了81.5%的ICL性能提升；4.实验使用了标准基准数据集，并提供了开源代码。

在最近的相关研究中，一些论文探索了Mamba SSM在特定任务上的性能，例如ICL，但是本论文通过在标准基准数据集上进行实验，证明了相比当前最先进的Transformer LLMs，Mamba SSM在ICL方面的性能提升仅达到了38%。