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

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【Mamba/扩散/多模态】交流群

添加微信：CVer5555，小助手会拉你进群！

扫描下方二维码，加入CVer学术星球！可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料，及最前沿应用！发论文/搞科研/涨薪，强烈推荐！

转载自：机器之心 | 编辑：蛋酱

LSTM：这次重生，我要夺回 Transformer 拿走的一切。

20 世纪 90 年代，长短时记忆（LSTM）方法引入了恒定误差选择轮盘和门控的核心思想。三十多年来，LSTM 经受住了时间的考验，并为众多深度学习的成功案例做出了贡献。然而，以可并行自注意力为核心 Transformer 横空出世之后，LSTM 自身所存在的局限性使其风光不再。

当人们都以为 Transformer 在语言模型领域稳坐江山的时候，LSTM 又杀回来了 —— 这次，是以 xLSTM 的身份。

5 月 8 日，LSTM 提出者和奠基者 Sepp Hochreiter 在 arXiv 上传了 xLSTM 的预印本论文。

论文的所属机构中还出现了一家叫做「NXAI」的公司，Sepp Hochreiter 表示：「借助 xLSTM，我们缩小了与现有最先进 LLM 的差距。借助 NXAI，我们已开始构建自己的欧洲 LLM。」

• 论文标题：xLSTM: Extended Long Short-Term Memory

• 论文链接：https://arxiv.org/pdf/2405.04517

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

(i) 无法修改存储决策。

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

(ii) 存储容量有限，即信息必须压缩成标量单元状态。

图 2 右侧给出了 Wikitext103 上不同 token 频率的 token 预测困惑度。由于 LSTM 的存储容量有限，它在不常见 token 上的表现较差。xLSTM 通过矩阵内存解决了这一问题。

(iii) 由于内存混合而缺乏可并行性，需要进行顺序处理。例如，从一个时间步到下一个时间步的隐藏状态之间的隐藏 - 隐藏连接。

与此同时，Sepp Hochreiter 和团队在这篇新论文中回答了一个关键问题：如果克服这些局限性并将 LSTM 扩展到当前大语言模型的规模时，能实现怎样的性能？

将 LSTM 扩展到数十亿参数

为了克服 LSTM 的局限性，xLSTM 对等式（1）中的 LSTM 理念进行了两项主要修改。

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

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

(i) sLSTM（第 2.2 节），具有标量内存、标量更新和内存混合功能；

(ii) mLSTM（第 2.3 节），具有矩阵内存和协方差（外积）更新规则，完全可并行处理。

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

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

xLSTM 块应在高维空间中对过去进行非线性总结，以便更好地分离不同的历史或上下文。分离历史是正确预测下一个序列元素（如下一个 token）的先决条件。研究者在此采用了 Cover 定理，该定理指出，在高维空间中，非线性嵌入模式比在原始空间中更有可能被线性分离。

他们考虑了两种残差块结构：(i) post up-projection 的残差块（如 Transformer），它非线性地概括了原始空间中的历史，然后线性地映射到高维空间，应用非线性激活函数，再线性地映射回原始空间（图 3 左侧和图 1 第三栏，更详细的版本见图 9）。(ii) pre up-projection 的残差块（如状态空间模型），它线性地映射到高维空间，在高维空间中非线性地总结历史，然后线性地映射回原始空间。对于包含 sLSTM 的 xLSTM 块，研究者主要使用了 post up-projection 块。对于包含 mLSTM 的 xLSTM 块，使用 pre up-projection 块，因为在高维空间中内存容量会变大。

实验

随后，研究者对 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 倍。

可以看出，xLSTM 在性能和扩展性上都更胜一筹。

更多研究细节，可参考原论文。

何恺明在MIT授课的课件PPT下载

在CVer公众号后台回复：何恺明，即可下载本课程的所有566页课件PPT！赶紧学起来！

CVPR 2024 论文和代码下载

在CVer公众号后台回复：CVPR2024，即可下载CVPR 2024论文和代码开源的论文合集

Mamba、多模态和扩散模型交流群成立

扫描下方二维码，或者添加微信：CVer5555，即可添加CVer小助手微信，便可申请加入CVer-Mamba、多模态学习或者扩散模型微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注：研究方向+地点+学校/公司+昵称（如Mamba、多模态学习或者扩散模型+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer5555，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集近万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看