xLSTM:扩展长短期记忆网络的综合指南

img

二十多年来, Sepp Hochreiter 开创性的 长短期记忆 (LSTM) 架构在众多深度学习的突破和实际应用中起到了重要作用。从生成自然语言到支持语音识别系统,LSTM 一直是推动 AI 革命的核心力量。

然而,即使是 LSTM 的创造者也认可其无法充分发挥全部潜力的固有局限性。这些局限性包括无法修正存储的信息、有限的存储能力以及缺乏并行处理功能,这些因素为变压器模型及其他模型的崛起铺平了道路,使其在更复杂的语言任务中超越了 LSTM。

但在最近的发展中,Hochreiter 及其在 NXAI 的团队推出了一种名为 扩展 LSTM (xLSTM) 的新变体,解决了这些长期存在的问题。在最近的一篇研究论文中提出的 xLSTM 建立在 LSTM 成功的基础上,通过架构创新克服了其关键弱点。

xLSTM 的核心是两个新颖的组件:指数门控和增强的记忆结构。指数门控允许更灵活地控制信息流动,使 xLSTM 能够在遇到新上下文时有效地修正决策。同时,引入的矩阵记忆与传统的标量 LSTM 相比,极大地增加了存储容量。

但改进不仅限于此。通过借用大型语言模型中的技术,如可并行化和残差块堆叠,xLSTM 可以有效地扩展到数十亿参数。这解锁了其对极长序列和上下文窗口进行建模的潜力——这一能力对于复杂的语言理解至关重要。

Hochreiter 最新创造的影响是巨大的。想象一下,虚拟助手可以可靠地追踪长达数小时的对话上下文;或者语言模型在广泛的数据训练后,可以更稳健地推广到新领域。应用范围涵盖了 LSTM 曾经影响的所有领域:聊天机器人、翻译、语音接口、程序分析等,但现在随着 xLSTM 的突破能力得到大幅提升。

在这篇深度技术指南中,我们将深入探讨 xLSTM 的架构细节,评估其新颖的组件如标量和矩阵 LSTM、指数门控机制、记忆结构等。您将从实验结果中了解 xLSTM 在性能上的显著提升,与当前最先进的架构如变压器和最新的递归模型相比。

理解起源:LSTM 的局限性

在我们深入探讨 xLSTM 的世界之前,了解传统 LSTM 架构所面临的局限性是至关重要的。这些局限性是驱动 xLSTM 和其他替代方法开发的重要因素。

  1. 无法修正存储决策:LSTM 的主要局限性之一是,当遇到更相似的向量时,它难以修正已存储的值。这可能导致在需要动态更新存储信息的任务中表现不尽如人意。

  2. 有限的存储容量:LSTM 将信息压缩到标量细胞状态中,这限制了它们在处理稀有标记或长程依赖时,有效存储和检索复杂数据模式的能力。

  3. 缺乏并行化能力:LSTM 中的记忆混合机制涉及时间步之间的隐隐连接,这强制了顺序处理,阻碍了计算的并行化,限制了扩展性。

这些局限性为变压器模型和其他架构的出现铺平了道路,使其在某些方面超过了 LSTM,特别是在扩展到更大模型时。

xLSTM 架构

扩展 LSTM (xLSTM) 家族

xLSTM 的核心是对传统 LSTM 框架的两项主要修改:指数门控和新颖的记忆结构。这些增强引入了两种新的 LSTM 变体,即 sLSTM (标量 LSTM)和 mLSTM (矩阵 LSTM)。

  1. sLSTM
    具有指数门控和记忆混合的标量 LSTM
    • 指数化门控:sLSTM 为输入和遗忘门引入了指数激活函数,从而实现对信息流更加灵活的控制。
    • 归一化与稳定性:为了防止数值不稳定,sLSTM 引入了一个归一化状态,该状态跟踪输入门和未来遗忘门的积。
    • 内存混合:sLSTM 支持多个内存单元,并允许通过递归连接进行内存混合,从而提取复杂模式和状态跟踪能力。
  2. mLSTM
    增强存储能力的矩阵 LSTM
    • 矩阵记忆:mLSTM 使用矩阵记忆代替标量记忆单元,增加了存储容量,并提高了信息检索效率。
    • 协方差更新规则:mLSTM 借鉴双向联想记忆(BAM)的灵感,采用协方差更新规则来高效地存储和检索键-值对。
    • 并行化:通过放弃内存混合,mLSTM 实现了完全并行化,使其在现代硬件加速器上能够高效计算。
      这两个变种,sLSTM 和 mLSTM,可以集成到残差块架构中,形成 xLSTM 块。通过残差堆叠这些 xLSTM 块,研究人员可以建立强大的 xLSTM 架构,以针对特定任务和应用领域。

数学原理

传统的 LSTM:

原始的 LSTM 架构引入了常数误差环和门控机制,以克服递归神经网络中的梯度消失问题。
The repeating module in an LSTM
LSTM 中的重复模块 – 来源
LSTM 记忆单元的更新由以下方程控制:
细胞状态更新:ct = ft ⊙ ct-1 + it ⊙ zt
隐藏状态更新:ht = ot ⊙ tanh(ct)
其中:

  • 𝑐𝑡 是时间 𝑡t 的细胞状态向量
  • 𝑓𝑡 是遗忘门向量
  • 𝑖𝑡 是输入门向量
  • 𝑜𝑡 是输出门向量
  • 𝑧𝑡 是由输入门调制的输入
  • ⊙ 表示按元素乘法
    门 ft、it 和 ot 控制从细胞状态 ct 存储、忘记和输出的信息,从而缓解梯度消失问题。

带有指数化门控的 xLSTM:

img
xLSTM 架构引入了指数化门控,允许对信息流进行更灵活的控制。对于标量 xLSTM (sLSTM) 变种:
细胞状态更新:ct = ft ⊙ ct-1 + it ⊙ zt
归一化状态更新:nt = ft ⊙ nt-1 + it
隐藏状态更新:ht = ot ⊙ (ct / nt)
输入和遗忘门:it = exp(W_i xt + R_i ht-1 + b_i) ft = σ(W_f xt + R_f ht-1 + b_f) OR ft = exp(W_f xt + R_f ht-1 + b_f)
输入门 (it) 和遗忘门 (ft) 的指数激活函数,与归一化状态 nt 一起,实现了对记忆更新和修订存储信息的更有效控制。

带有矩阵记忆的 xLSTM:

对于带有增强存储容量的矩阵 xLSTM (mLSTM) 变种:
细胞状态更新:Ct = ft ⊙ Ct-1 + it ⊙ (vt kt^T)
归一化状态更新:nt = ft ⊙ nt-1 + it ⊙ kt
隐藏状态更新:ht = ot ⊙ (Ct qt / max(qt^T nt, 1))
其中:

  • 𝐶𝑡 是矩阵细胞状态
  • 𝑣𝑡 和 𝑘𝑡 是值和键向量
  • 𝑞𝑡 是用于检索的查询向量
    这些关键方程突显 xLSTM 如何通过指数化门控实现更灵活的记忆控制,并通过矩阵记忆提高存储能力。这些创新的结合使 xLSTM 克服了传统 LSTM 的局限性。

xLSTM 的关键特性和优势

  1. 修订存储决策的能力:得益于指数化门控,xLSTM 可以在遇到更相关的信息时有效修订存储值,克服传统 LSTM 的重大限制。

  2. 增强的存储容量:mLSTM 中的矩阵记忆提供了增加的存储容量,使 xLSTM 能够更有效地处理稀有标记、长距离依赖和复杂数据模式。

  3. 并行化:xLSTM 的 mLSTM 变种完全可并行化,使其能够在现代硬件加速器(例如 GPU )上高效计算,并能够扩展到更大的模型。

  4. 内存混合和状态跟踪:xLSTM 的 sLSTM 变种保留了传统 LSTM 的内存混合功能,支持状态跟踪,使 xLSTM 在某些任务中比 Transformers 和 State Space Models 更具表达能力。

  5. 可扩展性:通过利用现代大型语言模型(LLMs)的最新技术,xLSTM 可以扩展到数十亿参数,从而在语言建模和序列处理任务中解锁新的可能性。

实验评估:展示 xLSTM 的能力

img

研究论文对 xLSTM 进行了全面的实验评估,重点展示了其在各种任务和基准测试中的表现。以下是一些关键发现:

  1. 综合任务和长范围竞技场

    • xLSTM 在需要状态跟踪的形式语言任务上表现出色,超越 Transformers、状态空间模型和其他 RNN 架构。
    • 在多查询关联记忆任务中,xLSTM 展现了增强的记忆能力,超过了非 Transformer 模型,并达到与 Transformers 相媲美的性能。
    • 在长范围竞技场基准测试中,xLSTM 一直表现稳定,展示了其处理长上下文问题的高效性。
  2. 语言建模和下游任务

    • 当在 SlimPajama 数据集的 150 亿个 token 上进行训练时,xLSTM 在验证困惑度方面超越了现有方法,包括 Transformers、状态空间模型和其他 RNN 变体。
    • 随着模型规模的扩大,xLSTM 继续保持其性能优势,表现出良好的扩展行为。
    • 在常识推理和问答等下游任务中,xLSTM 无论是在各种模型规模上,都优于最先进的方法。
  3. PALOMA 语言任务的表现

    • 在 PALOMA 语言基准测试的 571 个文本域中,xLSTM[1:0] (sLSTM 变体)在 99.5% 的域中比 Mamba 方法的困惑度低,比 Llama 方法低 85.1%,比 RWKV-4 方法低 99.8%。
  4. 扩展规律和长度外推

    • 在加在 3000 亿个 SlimPajama token 上训练时,xLSTM 表现出良好的扩展规律,表明其随着模型规模的增加,性能仍有进一步提升的潜力。
    • 在序列长度外推实验中,即使对于显著超过训练期间见过的上下文长度,xLSTM 模型也能保持低困惑度,超越其他方法。

这些实验结果突出了 xLSTM 的显著能力,使其成为语言建模任务、序列处理和广泛应用的有前途的竞争者。

实际应用和未来方向

xLSTM 的潜在应用涵盖从自然语言处理和生成到序列建模、时间序列分析等多个领域。以下是一些 xLSTM 可能产生重大影响的令人兴奋的领域:

  1. 语言建模和文本生成:凭借其增强的存储能力和修订存储信息的能力,xLSTM 可以革新语言建模和文本生成任务,实现更连贯、语境感知和流利的文本生成。

  2. 机器翻译:xLSTM 的状态跟踪能力在机器翻译任务中可能非常宝贵,在这些任务中,维持上下文信息和理解长距离依赖性对于准确翻译至关重要。

  3. 语音识别和生成:xLSTM 的并行化和可扩展性使其非常适合语音识别和生成应用,在这些应用中,高效处理长序列是必不可少的。

  4. 时间序列分析和预测:xLSTM 处理长距离依赖和有效存储提取复杂模式的能力,可能会显著改进跨多个领域如金融、天气预报和工业应用的时间序列分析和预测任务。

  5. 强化学习和控制系统:xLSTM 在强化学习和控制系统中的潜力非常大,因为其增强的记忆能力和状态跟踪能力可以在复杂环境中实现更智能的决策和控制。

架构优化和超参数调整

尽管当前的结果令人鼓舞,xLSTM 的架构仍有优化的空间和超参数调优的余地。研究人员可以探索不同的 sLSTM 和 mLSTM 块组合,在整体架构中变化比例和放置位置。此外,系统性的超参数搜索可能会带来进一步的性能改进,尤其是对于更大的模型。
img

硬件感知优化:为了充分利用 xLSTM 尤其是 mLSTM 变体的并行性,研究人员可以探讨针对特定 GPU 架构或其他加速器的硬件感知优化。这可能涉及优化 CUDA 核心、内存管理策略,以及利用专门的指令或库进行高效矩阵操作。

与其他神经网络组件的整合:探索将 xLSTM 与其他神经网络组件(例如注意力机制、卷积、或自监督学习技术)结合起来,可能会导致混合架构的出现,结合不同方法的优势。这些混合模型有可能解锁新能力,并在更广泛的任务中提升性能。

少样本学习和迁移学习:探索在少样本和迁移学习场景中使用 xLSTM 可能是未来研究的一个激动人心的方向。通过利用其增强的记忆能力和状态追踪能力,xLSTM 有可能实现更高效的知识转移,并在有限的训练数据下迅速适应新任务或领域。

可解释性和解释性:与许多深度学习模型一样,xLSTM 的内部运作可能是不透明且难以理解的。开发解释和说明 xLSTM 决策过程的方法,可以使模型更加透明和可信,促进其在关键应用中的采用,并提高其问责性。

高效且可扩展的训练策略:随着模型的规模和复杂性不断增长,高效且可扩展的训练策略变得越来越重要。研究人员可以探索专门为 xLSTM 架构设计的模型并行化、数据并行化和分布式训练方法,以便训练更大规模的模型,可能还会减少计算成本。

这些是一些关于 xLSTM 未来研究方向和进一步探索领域的潜在建议。

结论

xLSTM 的引入标志着在追求更强大、更高效的语言建模和序列处理架构方面的重要里程碑。通过解决传统 LSTM 的局限性,并利用指数门控和矩阵存储结构等新技术,xLSTM 在广泛的任务和基准测试中表现出了显著的性能。

然而,这一旅程并未就此终止。正如任何具有突破性的技术一样,xLSTM 为进一步探索、改进和在现实场景中的应用提供了令人兴奋的机会。随着研究人员不断突破可能性的边界,我们预计在自然语言处理和人工智能领域会看到更为惊人的进步。

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐