比较 Attention、Mamba、Jamba 和 xLSTM

序幕

我们在本文中只讨论时间复杂性——这是故意的。

有关空间复杂度，请参阅我关于 1 位变压器的文章，网址为：

HACKERNOON

Why 1-Bit Transformers Will Change the World | HackerNoon

Why are 1-bit transformers such a big deal? Read this article to find out!

介绍

就生成式人工智能技术而言，我们正在向未来迈进，大型语言模型背后的算法也不例外。在本文中，我们将介绍生成式人工智能领域最近最令人兴奋的三个发展，并详细讨论它们。其中之一还实现了运行大型语言模型算法的最佳时间复杂度。换句话说，最近的一项发展已成为最优化、最快的 LLM 变换算法——就渐近时间复杂度而言，按照我们目前的模型，除非通过常数时间优化，否则不可能比这更快。因为我们要处理数千亿个参数，所以常数的加速可能会相当大！我希望你和我一样兴奋，因为这将是一次激动人心的旅程！

现有算法——基于注意力机制的 Transformer

每个人都熟悉 2017 年具有开创性的《注意力就是你需要的一切》论文，但我还是要对其进行总结，以便新手能够更清楚地了解我们正在谈论的内容。

这是研究论文的链接：

你只需要注意力

ARXIV

[1706.03762] Attention Is All You Need

来自论文介绍：

循环神经网络，特别是长短期记忆和门控循环神经网络，已经成为序列建模和传导问题（例如语言建模和机器翻译）中最先进的方法。

从那时起，人们不断努力突破循环语言模型和编码器-解码器架构的界限。

循环模型通常根据输入和输出序列的符号位置进行计算。

通过将位置与计算时间中的步骤对齐，它们会生成一系列隐藏状态 ℎ𝑡，作为前一个隐藏状态 ℎ𝑡−1 和位置 𝑡 的输入的函数。

这种固有的顺序性阻碍了训练示例内的并行化，这在较长的序列长度下变得至关重要，因为内存限制限制了跨示例的批处理。

最近的研究通过分解技巧和条件计算显著提高了计算效率，同时在后者的情况下也提高了模型性能。

然而，顺序计算的基本约束仍然存在。

注意力机制已经成为各种任务中引人注目的序列建模和传导模型不可或缺的一部分，它允许对依赖关系进行建模，而无需考虑它们在输入或输出序列中的距离。

然而，除了少数情况外，这种注意力机制都是与循环网络结合使用的。

在这项工作中，我们提出了 Transformer，这是一种避免重复的模型架构，而是完全依靠注意力机制来绘制输入和输出之间的全局依赖关系。

Transformer 可以实现更加显著的并行化，在八块 P100 GPU 上仅训练 12 小时后，就能达到翻译质量的新水平。

众所周知，GPT-1、GPT-2、GPT-3 和 GPT 3.5 Transformer 很快就彻底改变了生成式人工智能。

突然间，机器可以讲看似人类的英语。

这是在接下来的两年里占据文章和研究新闻简报的经典图表：

开创性的 Transformer 架构。

然后 GPT-4 问世了——生活从此不再一样了。

我们已经跨过了一个临界点。

但是，这些变压器价格昂贵、训练缓慢，并且由于运营成本极高而难以部署。

Transformer 算法的时间复杂度是二次方，即O(n*n)，其中 n 是输入参数的数量。

对于具有 𝐿 层的标准 Transformer 模型，推理算法的时间复杂度为 𝑂( L*n*n*d ) ，其中 L 为层数，n 为输入标记数，d 为 Transformer 的深度。

暂时而言，这似乎是最先进的技术。

早在 2021 年，另一篇论文就引入了量化，这似乎是下一个最先进的机制（参见前言部分）。

但很快我们就将面临另一个竞争者。

欢迎来到 Mamba 算法

这是相关的研究论文：

Mamba：具有选择性状态空间的线性时间序列建模

ARXIV

[2312.00752] Mamba: Linear-Time Sequence Modeling with Selective State Spaces

摘自研究论文摘要：

基础模型目前为深度学习中大多数令人兴奋的应用提供支持，它们几乎普遍基于 Transformer 架构及其核心注意力模块。

许多亚二次时间架构（例如线性注意力、门控卷积和循环模型以及结构化状态空间模型 (SSM)）已被开发用于解决 Transformers 在长序列上的计算效率低下的问题，但它们在语言等重要模态上的表现不如注意力。

我们发现此类模型的一个主要弱点是无法进行基于内容的推理，并做出了一些改进。

首先，简单地让 SSM 参数成为输入的函数来解决它们使用离散模态的弱点，允许模型根据当前标记沿序列长度维度有选择地传播或忘记信息。

其次，尽管这种变化阻碍了有效卷积的使用，但我们还是设计了一种循环模式下的硬件感知并行算法。

我们将这些选择性 SSM 集成到简化的端到端神经网络架构中，无需注意力甚至 MLP 块（Mamba）。

Mamba 具有快速推理（吞吐量比 Transformers 高 5 倍）和序列长度的线性缩放功能，并且其性能在实际数据上可提升至数百万长度的序列。

作为通用序列模型主干，Mamba 在语言、音频和基因组学等多种模式下实现了最先进的性能。

在语言建模方面，我们的 Mamba-3B 模型在预训练和下游评估方面均优于同等大小的 Transformers，并匹敌两倍大小的 Transformers。

苏突然

突然间，我们城里有了一个新的竞争对手！

Mamba-Transformer 算法的主要优点是：

1. 混合架构：

   Mamba 结合了 Transformer 和 Mamba 层，从而提高了性能和吞吐量，同时保持了可管理的内存占用。

   
   

2. 面临的挑战：

   • 内存和计算效率：Mamba 解决了 Transformers 的高内存和计算要求，特别是对于长上下文。

   • 摘要状态：与 Transformers 不同，Mamba 提供了单一摘要状态，从而实现了更快的推理。

     
     

3. 建筑细节：

   • Transformer 层：Mamba 采用了 Transformer 层。
   • Mamba 层：Mamba 引入了状态空间建模 (SSM) 层，可以有效地处理长上下文。
   • 混合专家 (MoE) ：Mamba 使用 MoE 层来增加模型容量，而不会显著增加计算要求。

   
   

4. 表现：

   Mamba 的性能与 Mixtral-8x7B 和 Llama-2 70B 等现有模型相当，同时支持高达 256K 个标记的上下文长度。

   
   

5. 效率：对于长上下文，Mamba 的吞吐量是 Mixtral-8x7B 的 3 倍，并且即使在较大上下文中也能装入单个 GPU。

   
   

Mamba 架构代表了大型语言模型的重大进步，结合了 Transformer 和 SSM 方法的优势。

然而，随着测试的继续，发现 Mamba 算法并不适合所有用例。

特别是，Mamba 算法在面对 IMDB 数据集时惨遭失败。

然而，该架构仍然是最先进的，并且被发现在视觉用例中非常有用。

你可以在这里看到 Python 的实现：

GITHUB

GitHub - alxndrTL/mamba.py: A simple and efficient Mamba implementation in pure PyTorch and MLX.

A simple and efficient Mamba implementation in pure PyTorch and MLX. - alxndrTL/mamba.py

这是对 Mamba 算法的出色解释，并提供了理论。

TOWARDS DATA SCIENCE

Mamba: SSM, Theory, and Implementation in Keras and TensorFlow | Towards Data Science

Understanding how SSMs and Mamba work, along with how to get started with implementing it in Keras and TensorFlow.

以下是 PyPI 中的标准 Mamba 实现：

PYPI

Client Challenge

Mamba 算法已经过时了，现在仍然是一个高度活跃的研究领域。继任者很快就问世了，但我们会把最好的留到最后。

我们将讨论下一个算法——xLSTM 算法

xLSTM 算法

您可以参考此处的研究论文：

xLSTM：扩展长短期记忆 - arXiv 。

ARXIV

[2405.04517] xLSTM: Extended Long Short-Term Memory

摘自研究论文摘要：

在 20 世纪 90 年代，恒定错误轮播和门控被引入作为长短期记忆 (LSTM) 的核心思想。

从那时起，LSTM 经受住了时间的考验，并为众多深度学习的成功案例做出了贡献，特别是它们构成了第一个大型语言模型 (LLM)。

然而，以可并行自注意力为核心的 Transformer 技术的出现，标志着一个新时代的到来，其在规模上超越了 LSTM。

我们现在提出一个简单的问题：当将 LSTM 扩展到数十亿个参数，利用现代 LLM 的最新技术，但减轻 LSTM 已知的局限性时，我们在语言建模方面能取得多大进展？

首先，我们引入指数门控并采用适当的规范化和稳定技术。

其次，我们修改LSTM记忆结构，得到：

（i）具有标量记忆、标量更新和新记忆混合的 sLSTM，

(ii) mLSTM 可以通过矩阵存储器和协方差更新规则完全并行化。

将这些 LSTM 扩展集成到残差块主干中会产生 xLSTM 块，然后将其残差堆叠到 xLSTM 架构中。

指数门控和修改后的内存结构增强了 xLSTM 功能，与最先进的 Transformers 和状态空间模型相比，无论是性能还是扩展方面，它都具有良好的表现。

xLSTM 是 LSTM 模型的新生力量

长短期记忆 (LSTM) 算法在当时非常有用并且取得了相当多的成功。

xLSTM 使用相同的模型，但架构完全不同。

这是主要的创新，研究论文中的这张图表对此进行了总结：

xLSTM 的主要优点是：

xLSTM 与 Transformer 算法相比的优势：

1. 处理长序列：

   • xLSTM 专门设计用于处理长序列，其门控机制可控制信息流。与传统 LSTM 相比，这使其能够更有效地捕获序列数据中的长期依赖关系。

     
     

   
   

   
   

2. 计算效率：

   • 对于某些任务来说，xLSTM 的计算效率可能更高，尤其是在处理较小的数据集或序列长度不是特别大时。

   • 另一方面，Transformer 由于其自注意力机制而需要大量的计算资源，该机制与序列长度成二次方关系。

     
     

   
   

   
   

3. 内存使用情况：

   • xLSTM 通常比 Transformers 需要更少的内存。

   • Transformers 中的自注意力机制需要存储大量的注意力矩阵，这可能会占用大量内存，尤其是对于长序列而言。

     
     

4. 训练稳定性：

   • xLSTM 在训练期间可能更加稳定，因为它具有循环特性和有助于缓解梯度消失问题的门控机制。

   • Transformer 虽然功能强大，但有时训练起来更具挑战性，可能需要仔细调整超参数和正则化技术。

     
     

5. 简单性和可解释性：

   • 与 Transformers 相比，xLSTM 模型更容易理解和解释。

   • LSTM 的循环特性使其更容易追踪网络中的信息流，而 Transformers 中的注意力机制可能更加抽象且更难解释。

     
     

6. 小型数据集上的性能：

   • xLSTM 在小数据集或标记数据有限时表现更佳。

   • Transformer 通常需要大量数据才能发挥其全部潜力，因此在数据可用性有限的场景中，xLSTM 是更好的选择

     。

7. 顺序数据：

   • 对于某些类型的序列数据，例如时间序列或某些类型的自然语言处理任务，xLSTM 可能由于其处理序列的固有设计而提供更好的性能。

     
     

然而，值得注意的是，Transformers 有其自身的优势，例如更好的并行化能力、在大数据集上的卓越性能以及在许多 NLP 任务中的最先进的结果。

xLSTM 和 Transformer 之间的选择必须基于当前任务的具体要求和限制。

你可以在此处看到 PyTorch 中 xLSTM 的实现：

GITHUB

GitHub - kyegomez/xLSTM: Implementation of xLSTM in Pytorch from the paper: "xLSTM: Extended Long Short-Term Memory"

Implementation of xLSTM in Pytorch from the paper: "xLSTM: Extended Long Short-Term Memory" - kyegomez/xLSTM

你可以在这里看到xLSTM的详细解释：

这很好地总结了其当前状况：

HTTPS://WWW.HANNOVERMESSE.DE

AI & Machine Learning: XLSTM: What we know, what is still missing

Albert Ortig and Prof. Dr. Sepp Hochreiter from NXAI reveal details about the XLSTM approach in the Industrial AI Podcast. We listened and summarized the most important passages. The full interview with all the technical details can be found in the Industrial AI Podcast presented by HANNOVER MESSE.

但是 Mamba 的继任者实现了圣杯——LLM 算法的最佳时间复杂度

Jamba——曼巴继任者，毫无疑问地击败了它！

该研究论文可以在这里找到：

Jamba：混合 Transformer-Mamba 语言模型

ARXIV

[2403.19887] Jamba: A Hybrid Transformer-Mamba Language Model

摘自研究论文摘要：

我们提出了 Jamba，一种基于新颖的混合 Transformer-Mamba 专家混合 (MoE) 架构的新型基础大型语言模型。

具体来说，Jamba 将 Transformer 和 Mamba 层的块交错起来，享受两个模型系列的优势。

在某些层中添加了 MoE，以增加模型容量，同时保持活动参数的使用可管理。

这种灵活的架构允许特定于资源和目标的配置。

在我们实施的特定配置中，我们最终得到了一个适合单个 80GB GPU 的强大模型。

Jamba 是大规模构建的，与原始 Transformers 相比，它具有高吞吐量和小内存占用，同时在标准语言模型基准和长上下文评估中具有最先进的性能。

值得注意的是，该模型对于高达 256K 个标记上下文长度都呈现出了很好的效果。

我们研究了各种架构决策，例如如何结合 Transformer 和 Mamba 层，以及如何混合专家，并表明其中一些对于大规模建模至关重要。

我们还描述了 Jamba 的训练和评估所揭示的这些架构的几个有趣特性，并计划从各种消融运行中发布检查点，以鼓励进一步探索这种新颖的架构。

我们根据宽松的许可，公开提供 Jamba 实现的权重。

该实现可以在 HuggingFace 存储库中找到：

模型： https://huggingface.co/ai21labs/Jamba-v0.1

1. 与其他型号的比较：

   
   

   • 经过一系列基准评估，Jamba 的性能与其他最先进的型号（如Mistral-8x7B 、 Llama-2 70B和Mixtral-8x7B）相当。

     
     

   • 值得注意的是，Jamba 支持高达256K 个 token的上下文长度，这是公开模型中最长的

     
     

   • 在Hellaswag 、 Arc Challenge 和PIQA 等任务中，Jamba 的表现优于Llama 2 、 Mixtral 8x7B 和Gemma 等模型。

     
     

总而言之，Jamba 的混合架构结合了 Transformers 和 Mamba 层的优势，从而具有令人印象深刻的性能和可扩展性。

要记住的关键图表是上面这篇研究论文中提出的图表：

Mamba 和 Transformer 模型的交错导致时间复杂度的惊人增加，这在下面的文章中有精彩的总结：

Mamba 和 Jamba — 简单解释

• 作者： Nimrita Koul

• 日期：2024 年 4 月 1 日

  
  

• 概括：

  • Jamba 是 AI21 推出的第一个基于 Mamba 的生产级大型语言模型。

    
    

  • 它结合了 Transformer 和 Mamba 架构的优势。

  • 
    

    • 基于 Transformer 的模型由于其训练时间复杂度为O(n²)而难以处理长序列。

      
      

    • Mamba 架构提供线性训练时间复杂度（ O（n） ）和恒定推理时间（ O（1） ）。

      
      

    • Mamba 基于状态空间模型 (SSM) 架构。

      
      

    • SSM 使用状态变量来描述使用微分或差分方程的系统。

      
      

    • Mamba 的混合方法旨在解决现有模型的局限性。

      
      

      
      

您可以在此处阅读全文：

Mamba 和 Jamba — 简单解释，作者 Nimrita Koul，发表于 Medium.com。

已达到最佳界限！

这里要注意的关键点是，对于训练，算法必须查看每个输入标记至少一次，时间复杂度为O（n）。

此外，任何 LLM 模型的推理最快速度都是 O(1)——恒定时间，与标记的长度无关（一项令人难以置信的成就）！

在 Jamba 算法中，这两个限制都已达到！

因此，在恒定时间的改进下 - 这可能仍然非常高（这些数字是数千亿）：

Jamba 已达到 Transformer 算法的时间复杂度的最佳界限！

在给定的系统条件下，除非引入新技术（量子计算，任何人），否则我们根本无法获得更快的渐近时间复杂度！

这是一个非常重要的成果！

Jamba 是迄今为止最有前途的 Transformer 算法！

A121实验室官方公告：

AI21

Introducing Jamba: AI21's Groundbreaking SSM-Transformer Model

Debuting the first production-grade Mamba-based model delivering best-in-class quality and performance.

Medium 上另一篇关于 Jamba 的好文章：

MEDIUM

Unveiling Jamba: The First Production-Grade Mamba-Based Model

Introduction

目前 Jamba 的最佳实现之一：

再次，HuggingFace Hub 的 Jamba 模型：

HUGGINGFACE

ai21labs/Jamba-v0.1 · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

结论

因此，Jamba 达到了现有系统下电流互感器算法所能达到的极限时间复杂度，达到常数级变化。重复一遍；常数可能非常大，因为它们的数量级是数千亿！然而，这仍然是一项重大成就。而且，这方面的研究没有限制，尤其是与 DPO（直接偏好优化）和量化相结合时——更多信息请参见结语。

截至目前，实际上没有任何限制！

结语：

有一个方面似乎没有人公开致力于解决。

Mamba、xLSTM 和 Jamba 模型可以量化到 1 位精度吗？

当然！

我迫不及待地想看到 Mamba 和 Jamba 量化为 1 位后的性能提升！或者 1.58 位 {-1, 0, 1 }。

再次，请参阅本文以了解更多详细信息：

https://hackernoon.com/why-1-bit-transformers-will-change-the-world

这项技术的未来将会无比令人兴奋！

愿在这一领域工作的快乐和激动永远伴随着您！

干杯!

激动人心的时刻即将来临！

参考：

除了上述文章中明确提到的之外：

1. Transformer-XL：超越固定长度上下文的注意力语言模型
   • Dai, Z., Yang, Z., Yang, Y., Carbonell, J., & Le, QV (2019)。Transformer -XL：超越固定长度上下文的注意力语言模型。
2. Longformer：长文档转换器
   • Beltagy, I.、Peters, ME 和 Cohan, A. (2020)。Longformer ：长文档转换器。
3. Reformer：高效的变压器
   • Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020)。Reformer ：高效的变压器。
4. Linformer：具有线性复杂度的自注意力机制
   • Wang, S.、Li, BZ、Khabsa, M.、Fang, H. 和 Ma, H. (2020)。Linformer ：具有线性复杂度的自注意力。
5. 状态空间模型：时间序列数据建模的通用框架
   • Durbin, J., & Koopman, SJ (2012)。状态空间模型：时间序列数据建模的通用框架。
6. S4：使用结构化状态空间进行序列建模
   • Gu, A.、Goel, K. 和 Re, C. (2021)。S4 ：使用结构化状态空间进行序列建模。
7. 大型语言模型的计算效率
   • 匿名作者。（未指定年份）。[关于大型语言模型的计算效率]（未提供 URL）。
8. 高效 Transformer：调查
   • Tay, Y., Dehghani, M., Bahri, D., & Metzler, D. (2020).高效变压器：一项调查。
9. 神经网络的量化和训练，以实现高效的整数算术推理
   • Jacob, B., Kligys, S., Chen, B., Zhu, M., Tang, M., Howard, A., ... & Adam, H. (2018)。用于高效整数算术推理的神经网络量化与训练。
10. Q-BERT：基于 Hessian 的 BERT 超低精度量化
    • Shen, S., Dong, Z., Ye, J., Ma, L., & Gholami, A. (2020). Q-BERT：基于 Hessian 的 BERT 超低精度量化
11. BERT：用于语言理解的深度双向 Transformer 预训练
    • Devlin, J.、Chang, MW、Lee, K. 和 Toutanova, K. (2018)。BERT ：用于语言理解的深度双向 Transformer 预训练。
12. GPT-3：语言模型是小样本学习者
    • Brown, TB、Mann, B.、Ryder, N.、Subbiah, M.、Kaplan, J.、Dhariwal, P. ... & Amodei, D. (2020)。GPT -3：语言模型是小样本学习者。
13. RoBERTa：一种稳健优化的 BERT 预训练方法
    • Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., ... & Stoyanov, V. (2019)。RoBERTa ：一种稳健优化的 BERT 预训练方法。
14. ALBERT：用于语言表征自监督学习的精简版 BERT
    • Lan, Z., Chen, M., Goodman, S., Gimpel, K., Sharma, P., & Soricut, R. (2019)。ALBERT ：用于语言表征自监督学习的精简版 BERT 。
15. T5：使用统一的文本到文本转换器探索迁移学习的极限
    • Raffel, C.、Shazeer, N.、Roberts, A.、Lee, K.、Narang, S.、Matena, M.，... & Liu, PJ (2019)。T5 ：使用统一的文本到文本转换器探索迁移学习的极限。
16. DistilBERT，BERT 的精简版：更小、更快、更便宜、更轻便
    • Sanh, V., Debut, L., Chaumond, J., & Wolf, T. (2019). DistilBERT，BERT 的精简版：更小、更快、更便宜、更轻。

别忘了空间复杂度！几乎同样重要！

序言和尾声

对于量化，这篇论文绝对值得一读：

1. 研究论文 - BitNet：为大型语言模型扩展 1 位变压器：

   • 摘要：大型语言模型的规模不断扩大，给部署带来了挑战，并引发了人们对高能耗对环境影响的担忧。在本文中，我们引入了BitNet ，这是一种专为大型语言模型设计的可扩展且稳定的1 位 Transformer 架构。具体来说，我们引入BitLinear作为nn.Linear层的替代品，从头开始训练 1 位权重。语言建模的实验结果表明，与最先进的 8 位量化方法和 FP16 Transformer 基线相比， BitNet 实现了具有竞争力的性能，同时大幅降低了内存占用和能耗。此外， BitNet 表现出类似于全精度 Transformers 的缩放规律，表明它有可能有效扩展到更大的语言模型，同时保持效率和性能优势。

     阅读完整的研究论文

     
     

     https://arxiv.org/abs/2310.11453

     
     

   
   

HuggingFace 上的模型如下：

1. Hugging Face 存储库 - BitNet b1.58-3B 量化：

   • 该存储库包含BitNet b1.58-3B模型的量化版本。

     虽然原始存储库展示了令人印象深刻的验证结果，但它模拟了 BitNet 的线性层，导致内存使用量与 FP16 模型相似。

     
     

     探索 Hugging Face 模型

     
     

     https://huggingface.co/kousw/bitnet_b1_58-3B_quantized

关于未来，有一件事：它将会令人兴奋！