xLSTM: Extended Long Short-Term Memory

📄 arXiv: 2405.04517v2 📥 PDF

作者: Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, Sepp Hochreiter

分类: cs.LG, cs.AI, stat.ML

发布日期: 2024-05-07 (更新: 2024-12-06)

备注: Code available at https://github.com/NX-AI/xlstm

💡 一句话要点

xLSTM：通过扩展LSTM结构，提升语言建模性能并实现可扩展性。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: LSTM 语言建模 深度学习 循环神经网络 指数门控 并行计算 大型语言模型

📋 核心要点

  1. 传统LSTM在扩展到大规模语言模型时存在局限性，例如难以并行化和门控机制效率较低。
  2. 论文提出了xLSTM，通过引入指数门控和改进的记忆结构（sLSTM和mLSTM）来解决这些问题。
  3. 实验结果表明，xLSTM在性能和可扩展性方面均优于最先进的Transformer和状态空间模型。

📝 摘要（中文）

长短期记忆网络（LSTM）在20世纪90年代引入了恒定误差循环和门控机制，并经受住了时间的考验，为众多深度学习的成功做出了贡献，尤其是在大型语言模型（LLM）领域。然而，以并行化自注意力为核心的Transformer技术的出现标志着一个新时代的开始，并在规模上超越了LSTM。本文提出了一个简单的问题：当将LSTM扩展到数十亿参数，利用现代LLM的最新技术，并减轻LSTM已知的局限性时，语言建模能达到什么程度？为此，本文首先引入了具有适当归一化和稳定技术的指数门控。其次，修改了LSTM的记忆结构，得到了：（i）具有标量记忆、标量更新和新的记忆混合的sLSTM，（ii）具有矩阵记忆和协方差更新规则的完全可并行化的mLSTM。将这些LSTM扩展集成到残差块骨干网络中，产生xLSTM块，然后将这些xLSTM块残差堆叠成xLSTM架构。指数门控和修改后的记忆结构提升了xLSTM的性能，使其在性能和可扩展性方面都优于最先进的Transformer和状态空间模型。

🔬 方法详解

问题定义：现有LSTM在扩展到大型语言模型时面临挑战，主要体现在两个方面：一是难以有效利用现代LLM中的并行化技术，二是传统的门控机制在处理长序列时效率较低，限制了模型的性能和可扩展性。

核心思路：论文的核心思路是通过改进LSTM的门控机制和记忆结构，使其能够更好地利用并行计算资源，并提高处理长序列的能力。具体来说，引入了指数门控来提升门控效率，并设计了两种新的记忆结构：sLSTM和mLSTM，分别采用标量和矩阵记忆，以适应不同的计算需求。

技术框架：xLSTM架构基于残差块，将改进的LSTM块（xLSTM块）进行残差堆叠。每个xLSTM块包含指数门控和新的记忆结构（sLSTM或mLSTM）。整个架构类似于Transformer的残差结构，但使用xLSTM块代替了自注意力模块。

关键创新：论文的关键创新在于以下几点：1) 引入了指数门控，相比传统sigmoid门控，能够更有效地控制信息的流动；2) 提出了sLSTM和mLSTM两种新的记忆结构，其中mLSTM通过矩阵记忆和协方差更新规则实现了完全并行化，这是与传统LSTM的本质区别。

关键设计：指数门控采用指数函数代替sigmoid函数，并引入了归一化和稳定技术，以防止梯度爆炸或消失。sLSTM使用标量记忆和标量更新，适用于计算资源有限的场景。mLSTM使用矩阵记忆，并通过协方差更新规则实现并行化，适用于大规模计算。损失函数和优化器等其他技术细节与标准LLM类似，但针对xLSTM的特性进行了调整。

🖼️ 关键图片

img_0

📊 实验亮点

论文通过实验证明，xLSTM在语言建模任务上取得了显著的性能提升，与最先进的Transformer和状态空间模型相比，在性能和可扩展性方面均表现出优势。具体的性能数据和对比基线在论文中进行了详细展示，表明xLSTM是一种有竞争力的替代方案。

🎯 应用场景

xLSTM的潜在应用领域包括自然语言处理、语音识别、时间序列预测等。其改进的性能和可扩展性使其能够应用于更大规模的数据集和更复杂的任务，例如机器翻译、文本生成和对话系统。未来，xLSTM有望成为构建下一代大型语言模型的重要基石。

📄 摘要（原文）

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.