LSTM卷土重来！xLSTM：一举超越Mamba、Transformer！

创作时间:

作者:

@小白创作中心

LSTM卷土重来！xLSTM：一举超越Mamba、Transformer！

引用

来源

https://www.cnblogs.com/wxkang/p/18191885

LSTM（长短时记忆网络）作为深度学习领域的重要模型，在Transformer横空出世后一度风光不再。然而，LSTM的提出者Sepp Hochreiter近日在arXiv上发布了xLSTM的预印本论文，通过引入指数门控和矩阵内存等创新，成功克服了LSTM原有的局限性，并在语言建模任务中取得了超越现有最先进LLM的性能。

xLSTM的核心创新

具体来说，xLSTM从三个层面解决了LSTM以往所存在的局限性：

无法修改存储决策：通过"最近邻搜索"问题来举例说明这一局限性：在给定参考向量的情况下，必须按顺序扫描序列，寻找最相似的向量，以便在序列末端提供其附加值。图2左侧显示了这项任务的均方误差。当发现更相似的向量时，LSTM难以修改存储的值，而新的xLSTM通过指数门控弥补了这一限制。

存储容量有限：图2右侧给出了Wikitext103上不同token频率的token预测困惑度。由于LSTM的存储容量有限，它在不常见token上的表现较差。xLSTM通过矩阵内存解决了这一问题。
缺乏可并行性：由于内存混合而缺乏可并行性，需要进行顺序处理。例如，从一个时间步到下一个时间步的隐藏状态之间的隐藏-隐藏连接。

xLSTM的架构设计

为了克服LSTM的局限性，xLSTM对等式（1）中的LSTM理念进行了两项主要修改。在原来的LSTM中，恒定误差选择轮盘是由单元输入z_t对单元状态c_(t-1)（绿色）进行的加法更新，并由sigmoid门（蓝色）进行调节。输入门i_t和遗忘门f_t控制这一更新，而输出门o_t控制存储单元的输出，即隐藏状态h_t。存储单元的状态被ψ归一化或压缩，然后输出门控得到隐藏状态。

xLSTM的修改包括指数门控和新颖的内存结构，因此丰富了LSTM家族的两个成员：

sLSTM（第2.2节）：具有标量内存、标量更新和内存混合功能；
mLSTM（第2.3节）：具有矩阵内存和协方差（外积）更新规则，完全可并行处理。

sLSTM和mLSTM都通过指数门控增强了LSTM。为了实现并行化，mLSTM放弃了内存混合，即隐藏-隐藏递归连接。mLSTM和sLSTM都可以扩展到多个存储单元，其中sLSTM具有跨单元内存混合的特点。此外，sLSTM可以有多个头，但不存在跨头的内存混合，而只存在每个头内单元间的内存混合。通过引入sLSTM头和指数门控，研究者建立了一种新的内存混合方式。对于mLSTM而言，多头和多单元是等价的。

将这些新的LSTM变体集成到残差块模块中，就得到了xLSTM块。将这些xLSTM块剩余堆叠到架构中，就形成了xLSTM架构。xLSTM架构及其组件见图1。

实验评估

研究者对xLSTM进行了实验评估，并将其与现有的语言建模方法进行了比较。第4.1节讨论了xLSTM在合成任务中的具体能力。首先，研究者测试了xLSTM的新指数门控与内存混合在形式化语言上的有效性。然后，他们评估了xLSTM的新矩阵内存在多次查询联想记忆任务（Multi-Query Associative Recall，MQAR）中的有效性。最后，研究者评估了xLSTM在Long Range Arena（LRA）中处理长序列的性能。

第4.2节比较了当前各种语言建模方法的验证集复杂度，包括在同一数据集上对xLSTM进行消融研究，然后对不同方法的缩放行为进行评估。研究者在自回归语言建模设置中使用SlimPajama的15B token训练了xLSTM、Transformers、状态空间模型（SSM）等模型。表1中的结果显示，xLSTM在验证复杂度方面优于所有现有方法。图6显示了该实验的扩展结果，表明xLSTM对于更大规模的模型也有良好的表现。消融研究则表明，性能改进源于指数门控和矩阵内存。

第4.3节进行了更深入的语言建模实验。研究者增加了训练数据量，对来自SlimPajama的300B个token进行了训练，并比较了xLSTM、RWKV-4、Llama和Mamba。他们训练了不同大小的模型（125M、350M、760M和1.3B），进行了深入的评估。首先，评估这些方法在推断较长语境时的表现；其次，通过验证易混度和下游任务的表现来测试这些方法；此外，在PALOMA语言基准数据集的571个文本域上评估了这些方法；最后，评估了不同方法的扩展行为，但使用的训练数据多了20倍。