LSTM再升级！原作者携xLSTM回归，扩展LSTM到数十亿参数

Ne0inhk

25 Dec 2024 — 10 min read

LSTM再升级！原作者携xLSTM回归，扩展LSTM到数十亿参数

TSer 2024年08月20日 20:33 浙江

时序人

专注于时间序列领域下的科研、工业干货分享，紧跟AI+等领域的科技前沿

长短期记忆网络（Long Short-Term Memory）是一种特殊的循环神经网络（RNN），它能够学习到长距离的依赖关系。LSTM 的设计初衷是为了解决传统 RNN 在处理长序列数据时出现的梯度消失或梯度爆炸问题。

LSTM 的问世不仅在理论上是一大步，还在实际应用中也产生了革命性的影响。它很快成为了处理序列数据的利器，特别是在自然语言处理领域，LSTM 几乎参与了所有重要的任务，如语音识别、机器翻译、文本生成和情感分析等。

今年上半年，LSTM 提出者和奠基者 Sepp Hochreiter 携团队再次提出了新的 LSTM —— xLSTM，将 LSTM 扩展到数十亿参数规模。与 Transformer 和 State Space Models 等最先进的方法相比，xLSTM 模型在语言建模方面表现良好。扩展定律表明，更大的 xLSTM 模型将成为当前使用 Transformer 技术构建的大型语言模型（Large Language Models）的有力竞争者。xLSTM 有可能对强化学习、时间序列预测或物理系统建模等其他深度学习领域产生重大影响。

本文来聊聊 xLSTM 具体做了哪些改进创新。

www.zeeklog.com - LSTM再升级！原作者携xLSTM回归，扩展LSTM到数十亿参数

【论文标题】xLSTM: Extended Long Short-Term Memory

【论文地址】https://arxiv.org/abs/2405.04517

【开源代码】https://github.com/NX-AI/xlstm

LSTM回顾与局限

在20世纪90年代，长短期记忆网络（Long Short-Term Memory, LSTM）的核心思想——恒定误差旋转木马和门控机制被引入，用以解决循环神经网络中的梯度消失问题。自那时以来，LSTM 已经经历了时间的考验，并为众多深度学习的成功案例做出了贡献，特别是它们构成了第一批大型语言模型（LLMs）。

原始 LSTM 的计算涉及几个关键组件：输入门、遗忘门、输出门和单元状态。

图1：LSTM 单元的标准计算步骤

虽然 LSTM 在处理序列数据方面取得了巨大成功，但它们面临三个主要的局限性，这些局限也是驱动 xLSTM 和其他替代方法开发的重要因素。

LSTM 在处理最近邻搜索问题时难以修订其存储决策，这限制了其在动态更新信息方面的能力。

LSTM 的存储容量有限，这导致它在处理稀有标记预测任务时表现不佳，因为信息必须被压缩进标量单元状态中。

由于隐藏状态之间的连接，LSTM 缺乏并行处理的能力，这影响了其在大规模数据处理中的效率。

xLSTM的创新与架构

xLSTM 的关键在于“扩展”这个词，它不仅仅是对 LSTM 的简单改进。xLSTM 将原始 LSTM 扩展为 LSTM 的多种变体，例如 sLSTM 和 mLSTM，每种变体都针对特定的性能和功能进行了优化，以处理各种复杂的序列数据问题。

总体而言，xLSTM 主要通过两方面改进来增强LSTM的能力：（1）引入指数门控；（2）引入新型记忆结构（sLSTM 和 mLSTM）

图2：xLSTM家族

指数门控

这是对 LSTM 中传统门控机制的一种改进，允许模型更有效地更新其内部状态。指数门控通过引入适当的归一化和稳定技术，使得 LSTM 能够更好地处理信息流，特别是在需要修订存储决策的场景中。

新型记忆结构

（1）sLSTM

sLSTM 在 LSTM 的基础上添加了标量更新机制。该设计通过对内部存储单元进行细粒度控制来优化门控机制，使其更适合处理具有细微时间变化的序列。

图3：sLSTM的前向传播方程

sLSTM 具备了跨单元的内存混合能力，允许不同存储单元之间进行信息交互。sLSTM 还支持多头部结构，每个头部可以独立处理信息，但头部之间并不直接混合内存，而是在每个头部内部的单元之间进行内存混合。这种设计为模型提供了一种新的内存混合方式，有助于捕捉更复杂的序列特征。

（2）mLSTM

mLSTM 通过将原始 LSTM 中的向量运算扩展为矩阵运算，显著增强了模型的内存容量和并行处理能力。

为了提高计算效率并实现并行处理，mLSTM 特别放弃了传统的隐藏状态间的递归连接，即内存混合。对于 mLSTM 而言，多头结构和多单元结构在功能上是等价的，这意味着它可以灵活地根据需要选择使用哪一种结构来优化性能和计算效率。

图4：mLSTM的前向传播方程

xLSTM架构

将 sLSTM 和 mLSTM 这两种新型 LSTM 变体融合进残差块中，构成了 xLSTM 块。这些 xLSTM 块通过剩余连接的方式堆叠起来，形成了完整的 xLSTM 架构。

图5：xLSTM blocks

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

他们考虑了两种残差块架构：

post up-projection 的残差块（如Transformer），它在原始空间中非线性地总结过去，然后线性映射到高维空间，应用非线性激活函数，并线性映射回原始空间。（见图5的左面板和图2的第三列）

pre up-projection 的残差块（如状态空间模型），它线性映射到高维空间，在高维空间中非线性地总结过去，然后线性映射回原始空间。对于包含 sLSTM 的 xLSTM 块，主要使用 post up-projection 块。对于包含 mLSTM 的 xLSTM 块，使用 pre up-projection 块，因为在高维空间中内存容量变得更大。

与 Transformer 不同，xLSTM 网络在计算上具有线性复杂度，在内存上具有与序列长度相关的恒定复杂度。由于 xLSTM 内存具有压缩性，因此非常适合工业应用和边缘部署。

xLSTM性能分析

研究者的实验结果表明，xLSTM 在多个任务上展现出了卓越的性能，包括在合成任务、长距离竞技场（Long Range Arena）以及大规模语言建模实验中的表现。这表明 xLSTM 能够有效地处理长序列和大型上下文，并且在语言建模方面具有强大的能力。

xLSTM 通过引入指数门控和矩阵内存结构，显著提高了模型的扩展性和内存容量。在多查询关联记忆（Multi-Query Associative Recall）任务中，xLSTM 表现出了优秀的记忆能力，尤其是在需要记忆大量键值对时。

图6：xLSTM的指数门控与记忆混合测试

xLSTM 在与 Transformers 和 State Space Models（状态空间模型）等现有技术的比较中表现出色。在 150 亿个 token 的训练数据上，xLSTM 实现了最低的验证集困惑度（perplexity），并且在 300 亿个 token 的训练数据上，xLSTM 在不同模型尺寸下仍然保持了优异的性能。据此研究人员认为 xLSTM 未来有望能够与 Transformer 进行 PK。

表1：在SlimPajama数据集上训练的不同语言模型的比较

总结

xLSTM 通过引入指数门控和矩阵内存结构，显著提升了 LSTM 的存储容量和并行计算能力，为大型语言模型和深度学习领域提供了一种强大的新型序列处理工具，具有推动自然语言处理和序列预测任务发展的学术意义。虽然 xLSTM 在多个方面表现出色，但仍存在一些局限性，主要涉及计算效率、内存容量、实现优化和泛化能力等方面。尽管存在这些挑战，xLSTM 在语言建模和其他序列建模任务上展现出的潜力表明，通过进一步的研究和开发，这些局限性有可能被克服或减轻。