了解xLSTM，它会威胁到Transformer的统治地位吗？

前言

本文属于搬运内容，原作者：Salvatore Raieli

文章末尾附有原文链接

xLSTM与Transformer

“艺术要么是剽窃，要么是革命。” —— 保罗·高更

最初是RNN（循环神经网络），它们曾经表现不错。一段时间后，大家开始抱怨它们，想要一些更好的东西，能够最终与语言分析的统计方法相媲美，并且让我们摆脱那些令人讨厌的正则表达式。

此时，长短期记忆（LSTM）悄然到来，也许是因为这种结构太复杂，我们在各种门控和向量中迷失了方向。霍赫赖特（Hochreiter）是对的；他的LSTM确实优于RNN，前提是你有足够的计算能力来训练它们。正因如此，曾经有人尝试寻找更轻量级的替代方案，比如GRU（门控循环单元）。

长短期记忆（LSTM）单元可以顺序处理数据，并且保持其隐状态。图像来源：此处

然后，Transformer和注意力机制席卷了一切，像雪崩一样，它接管了所有的RNN及其变体。后来，它开始入侵原本不属于自己的领域，如计算机视觉和生物信息学。

Transformer的挽歌？

Transformer是否会成为引领我们走向人工通用智能的模型？还是会被取代？

一段时间以来，大家谈论的只有Transformer，适用于所有应用。作为ChatGPT和大规模语言模型（LLMs）成功的帮凶，Transformer似乎毫无争议地占据了主导地位。然而，一年前，似乎发生了一些变化，首先是Hyena的推出，然后是Mamba、Megalodon等，接踵而至。

如今，霍赫赖特已经将LSTM进行了并行化。

它们是什么？发生了什么？

LSTM（长短期记忆网络）已成功应用于多个领域，并且在文本生成任务中曾占据主导地位，直到2017年Transformer的出现（来源）。

尽管LSTM结构复杂，但它在生成任务中表现优秀，如翻译、编写代码、生成图像标题等。然而，Transformer迄今为止在这些任务上表现更好，以至于目前用于这些任务的模型都是基于Transformer的。

LSTM简要回顾：

LSTM由不同的元素组成，包括一个在时间步之间传递信息的旋转木马上，以及控制信息写入或遗忘的门控。最后，输出门控制将信息传递到下一层。

LSTM的三大问题：

• 无法修正存储决策。这主要是由于其顺序性质所致。Transformer在这方面也有局限性，但为大规模语言模型（LLM）适应的技术已减少了这一问题。

• 存储容量有限。这是因为部分信息必须被压缩为标量向量。Transformer通过多个注意力头（attention heads）不仅能够存储知识，还能建模这些知识之间的关系。

• 由于内存混合，无法并行化。Transformer的优势之一是其高度可并行化，这一特性正是在GPU普及的时代，使得Transformer能够在并行处理方面大放异彩。

xLSTM是什么？

xLSTM不仅仅是一个技术突破，它是朝着一个未来迈进的一步——在语言处理的效率、准确性和理解力方面，可以达到甚至超越人类能力。—— Sepp Hochreiter

作者通过提出两种可组合的模块，尝试解决LSTM的这些问题。

对于作者来说，LSTM无法修正存储决策的问题源于sigmoid函数及其遗忘门（forget gate）。他们用指数门控（exponential gating）替换了sigmoid函数，这种方法可以更动态地管理信息流。实际上，sigmoid函数呈S型，这使得处理高输入值时会遇到问题。如果重要输入值出现在这个曲线的部分，就很难判断是否应该保留这些信息。而指数函数解决了这个问题。显然，较大的动态范围会带来稳定性的问题，但这个问题可以通过先进的归一化技术来抵消。实际上，这种归一化是通过一种类似于滑动平均的方法来实现的。

然而，这可能仍然不够。指数函数有一个坏习惯，那就是它会迅速增长到极大的值并导致溢出。因此，作者增加了一个稳定器（简单来说，就是一个对数项）。稳定器取遗忘门和输入门的对数最大值，并从它们中减去这个值，以此来稳定这些值。

另一个重要的问题是并行化。没有并行化，使用训练LLM的硬件训练LSTM将是一场噩梦。问题的根源是细胞状态（cell state），它是一个标量，作者将其替换为矩阵（GPU的优势正是用来进行矩阵乘法的！）。实际上，由于这些矩阵是可并行化的，它们可以有多个头，就相当于拥有多头结构。这些结构在数据存储和检索上也更具优势，因此优化了大规模数据集的处理。

一旦我们拥有了这些元素，就可以将它们组合成一个xLSTM模块。作者提出了两个模块。

一个残差sLSTM块。由sLSTM残差连接、可选的卷积层和门控多层感知机（gated MLP）组成。这个块可以将数据映射到更高维的空间，然后再映射回原始空间。

一个残差mLSTM块。在这种情况下，我们将数据投影到一个更高的维度空间，进行卷积步骤以获取查询、键和值，作为mLSTM的输入，最后再将数据带回原始投影空间。同样，这里也有残差连接。

我们可以将这些层堆叠起来，创建一个xLSTM网络。

接下来，我们需要做一些考虑。显然，作者已经学到了Transformer的经验教训，理解了它的优点和缺点：

• 残差连接使得非常深的模型可以进行训练。残差连接也可以部分学习，从而进一步调整。

• 归一化步骤用于稳定训练。注意，关键输入通过1/√d（输入大小）进行归一化，这与自注意力机制中的做法一致。

• 多头机制的思想确保了更强的表达能力。

• 该模型具有线性计算复杂度，并且在序列长度上具有恒定的内存复杂度（而自注意力机制的复杂度是二次的）。

此时，作者决定在300B个tokens上训练他们的模型，并将其与其他架构（包括Transformer、状态空间模型和递归神经网络）进行比较，使用相同的数据集。他们还实验了sLSTM和mLSTM的组合以及比例。

我们可以推测，xLSTM在需要进行记忆混合或状态跟踪（如奇偶任务）来解决问题的任务中具有优势。在这种情况下，Transformer或状态空间模型之所以无法精确解决问题，是因为它们没有进行状态跟踪。

通常，在使用相同数据集和相同参数数量进行训练时，xLSTM模型与其他架构具有竞争力。

消融研究表明，所有新增的元素对于实现这种性能都是必不可少的。

正如讨论中提到的，这些门的作用非常重要，作者还对这些门进行了选择性消融，显示它们对性能有增量的正面影响。然而，结果表明，这些门的作用并非完全不可或缺。

显然，作者还测试了上下文长度（毕竟，这是当前的热门话题）。xLSTM模型在较长上下文的困惑度表现较低，明显优于其他模型。

最后，他们测试了下游任务，结果显示，xLSTM模型在所有模型规模中，在验证集困惑度上表现最佳（除了少数情况下Mamba似乎更好，但作者没有进一步研究）。

xLSTM的局限性

根据作者的观点，仍然存在一些局限性：

• sLSTM 阻碍并行化。作者通过开发一个快速的 CUDA 内核来解决这一问题，使得 sLSTM 块的速度比 mLSTM 慢仅 1.5 倍。

• mLSTM 的 CUDA 内核尚未优化，因此它在 FlashAttention 中表现不佳。

• 仍然存在复杂度为 d × d 的矩阵，因此通过扩展上下文，也会增加矩阵的大小。这些大矩阵可能会影响模型的高效使用。

• 架构或超参数优化仍然不足（例如，忘记门的初始化需要小心选择）。

我们部分回答了一个简单的问题：当将 LSTM 扩展到数十亿参数时，语言建模能达到什么程度？截至目前的回答是：“至少能达到当前技术，如 Transformer 或状态空间模型的水平”。

除了对 LSTM 扩展和并行化的理论兴趣，xLSTM 的潜在影响是什么呢？

结果看起来并不那么有前景。LSTM 受益于大量的 NLP 进展，但其中许多进展是受到 Transformer 启发的。这些进展不足以说服社区放弃庞大的生态系统，这些生态系统已经为 Transformer 训练和使用进行了优化。

xLSTM 的真正优势在于 状态跟踪。如果过去几年教给我们一件事，那就是 上下文学习 是一个非常强大的概念。我们无需重新训练模型，就可以创建技术来克服 Transformer 在解决任务时的局限性。因此，这一局限性可以通过一些新的提示技术来克服。

尽管如此，xLSTM 在 30 年后才出现仍然很有趣。也许，几年前它可能能够与 Transformer 竞争，但今天这篇论文的结果看起来像是一次 惨胜。它们已经失去了这个“王国”。无论如何，今天新架构的涌现意味着我们希望超越 Transformer。社区正在寻找一种新的架构，能够带我们走向 人工通用智能。这不会是 xLSTM，也不会是 Transformer。

参考

Hochreiter, 1997, Long Short-Term Memory

Beck, 2024, xLSTM: Extended Long Short-Term Memory

Vaswani, 2017, Attention Is All You Need

非官方代码 xLSTM

原文链接