扩展长短期记忆 xLSTM

2024 年 5 月 8 日,长短期记忆 LSTM 提出者和奠基者 Sepp Hochreiter  在 arXiv 上传了 xLSTM 的预印本论文 「xLSTM: Extended Long Short-Term Memory」。文章提出了一个问题:利用 LLM 的最新技术,将 LSTM 扩展到数十亿个参数时,语言建模方面能走多远?本文介绍了 LSTM 设计的重大进步,解决了传统 LSTM 的局限性,并引入了新功能来增强其在大型语言模型 (LLM) 中的性能。 

xLSTM 全称 Extended Long Short-Term Memory,xLSTM 复兴了长短期记忆 (LSTM) 思想,即恒定误差轮播和门控的概念。由 Sepp Hochreiter 和 Jürgen Schmidhuber 推出的 LSTM 是 20 世纪 90 年代革命性的深度学习架构,它成功克服了时序任务(例如时间序列或语言建模)的梯度消失问题。从那时起,LSTM 经受住了时间的考验,并为众多深度学习成功案例做出了贡献,特别是它们构成了第一个大型语言模型 (LLM) 。然而,以并行自注意力为核心的 Transformer 技术的出现标志着一个新时代的到来,在规模上超越了 LSTM 。

xLSTM 系列及其组件简介

如上图所示,概述了 xLSTM 系列及其组件。从左到右:

  1. 具有恒定误差轮播和门控的原始 LSTM 存储单元。
  2. 引入了两个新的存储单元:
  • sLSTM(标量 LSTM)具有指数门控和新的内存混合技术。
  • mLSTM(矩阵 LSTM)具有指数门控、并行训练、协方差更新规则和矩阵存储单元状态。

3. 将 mLSTM 和 sLSTM 存储单元集成到残差块中以形成 xLSTM 块。

4. xLSTM 架构是通过残差堆叠 xLSTM 块构建的。

xLSTM 对大语言模型 (LLM) 的意义

xLSTM 架构的引入对大语言模型 (LLM) 的开发和性能具有重大影响。通过解决传统 LSTM 的局限性并结合指数门控、矩阵存储器和可并行架构等新颖组件,xLSTM 为 LLM 开辟了新的可能性。

xLSTM 对于大语言模型 (LLM) 的主要优势之一是能够有效处理长序列和大规模语言建模任务。 xLSTM 的线性时间复杂度和恒定的内存复杂度使其非常适合处理冗长的文本数据,而不会导致与基于 Transformer 的模型相关的计算成本和内存使用量呈二次方增加。这种效率优势对于 LLM 来说尤其有价值,因为 LLM 通常需要在训练和推理过程中处理大量文本数据。

此外,与 Transformer LLM 和 RWKV 相比,xLSTM 的语言建模性能有所提高,其困惑度得分较低,表明它有可能提高 LLM 中生成文本的质量和连贯性。 xLSTM 中的矩阵记忆和指数门控机制使其能够从训练数据中捕获和保留更全面、更细致的信息,从而获得更好的语言理解和生成能力。

xLSTM 论文中提出的缩放定律表明,即使在更大的数据集(例如 300B token 的 SlimPajama 语料库)上进行训练,xLSTM 的性能优势仍然存在。这种可扩展性对于 LLM 来说至关重要,因为他们通常依赖大量的训练数据来实现最先进的性能。 xLSTM 在更大范围内保持其效率和建模能力的能力使其成为未来 LLM 的有前景的架构。

此外,xLSTM 架构的灵活性允许不同比例的 mLSTM 和 sLSTM 模块,为定制和适应特定语言建模任务提供了机会。这种适应性对于 LLM 来说很有价值,因为它们通常应用于具有不同要求和特征的各种自然语言处理任务。

xLSTM 架构还为 LLM 的研究和创新开辟了新途径。 xLSTM 中指数门控和矩阵内存的引入挑战了基于 Transformer 的模型的主导地位,并鼓励探索可提供更高效率和性能的替代架构。 xLSTM 的成功可能会激发对 LLM 的新型内存结构、门控机制和并行化技术的进一步研究。

总之,xLSTM 架构为 LLM 带来了重大进步。它的效率、可扩展性和改进的语言建模功能使其成为基于 Transformer 的模型的有前途的替代品。随着 LLM 领域的不断发展,xLSTM 引入的见解和创新可能会塑造未来的发展,并突破自然语言处理的可能性界限。 xLSTM 论文为 LLM 的新时代奠定了基础,可以有效处理大量文本数据,同时提供高质量的语言理解和生成。

参考来源

【1】xLSTM: Enhancing Long Short-Term Memory for Large Language Models