[论文解读] xLSTM: Extended Long Short-Term Memory

Maximilian Beck, Korbinian Pöppel+7 more|arXiv (Cornell University)|May 7, 2024

Topic Modeling被引用 86

一句话总结

xLSTM 引入指数门控和两个记忆增强变体（sLSTM 使用标量记忆，mLSTM 使用矩阵记忆），以克服 LSTM 的局限，在语言建模性能和可扩展性方面比 Transformers 和 State Space Models 有利。

ABSTRACT

In the 1990s, the constant error carousel and gating were introduced as the central ideas of the Long Short-Term Memory (LSTM). Since then, LSTMs have stood the test of time and contributed to numerous deep learning success stories, in particular they constituted the first Large Language Models (LLMs). However, the advent of the Transformer technology with parallelizable self-attention at its core marked the dawn of a new era, outpacing LSTMs at scale. We now raise a simple question: How far do we get in language modeling when scaling LSTMs to billions of parameters, leveraging the latest techniques from modern LLMs, but mitigating known limitations of LSTMs? Firstly, we introduce exponential gating with appropriate normalization and stabilization techniques. Secondly, we modify the LSTM memory structure, obtaining: (i) sLSTM with a scalar memory, a scalar update, and new memory mixing, (ii) mLSTM that is fully parallelizable with a matrix memory and a covariance update rule. Integrating these LSTM extensions into residual block backbones yields xLSTM blocks that are then residually stacked into xLSTM architectures. Exponential gating and modified memory structures boost xLSTM capabilities to perform favorably when compared to state-of-the-art Transformers and State Space Models, both in performance and scaling.

研究动机与目标

• 在大规模语言建模中克服经典 LSTM 限制（存储的修订、有限记忆、缺乏并行性）的需求动机。
• 提出两个记忆增强的 LSTM 变体（sLSTM 和 mLSTM）带有指数门控。
• 引入将这些变体集成到残差结构中的 xLSTM 块，以实现可扩展建模。
• 通过合成任务和大规模语言建模实验，证明在某些设置下 xLSTM 能与 Transformer/SSM 基线匹配甚至超过。

提出的方法

• 引入带归一化和稳定化的指数门控，以实现可修订的存储决策。
• 开发具有标量记忆、标量更新和跨单元/头的记忆混合的 sLSTM。
• 开发具有矩阵记忆和协方差更新规则的 mLSTM，以实现高容量存储和可并行检索。
• 通过将 sLSTM 或 mLSTM 嵌入残差块骨干（在变体取决于预上投影/后上投影）来创建 xLSTM 块。
• 将 xLSTM 块堆叠成 xLSTM 架构，使用 pre-LayerNorm 的骨干实现可扩展建模。
• 提供 CUDA 优化实现，并讨论相对于 Transformer 的内存/计算特性。

实验结果

研究问题

• RQ1 在规模化情况下，指数门控和新的记忆结构是否能修复 LSTM 的基本局限（存储再审、记忆容量和序列瓶颈）？
• RQ2 当参数达到十亿级规模时，sLSTM 和 mLSTM 在语言建模任务中与 Transformer 和 State Space Models 的比较如何？
• RQ3 带有残差堆叠的 xLSTM 架构是否在模型规模和数据规模下提供具有竞争力的困惑度和下游任务性能？

主要发现

• 指数门控和矩阵/标量记忆设计使 xLSTM 在验证困惑度上的语言建模任务中优于若干基线模型。
• 在合成与长距离任务中，xLSTM 展现出比传统 LSTM 以及某些 Transformer/SSM 变体更好的状态跟踪和记忆容量。
• xLSTM[1:0] 与 xLSTM[7:1] 在 SlimPajama 实验中对不同模型规模实现强劲的验证困惑度，并呈现有利的扩展行为。
• 在 300B-token 训练中，xLSTM 变体保持更好的序列长度外推、较长上下文困惑度以及与 RWKV、Llama、Mamba 相当的下游任务性能。
• 消融研究表明指数门控和矩阵记忆是性能提升的主要驱动因素，可学习的、输入相关的门控提供额外收益。
• xLSTM 架构在具有线性记忆考虑和可并行组件的情况下能够提供有竞争力的性能，作为 Transformer 主导的区间的可扩展替代方案。