xLSTM : 一个关于扩展长短期记忆的综合指南

理解起源：LSTM 的局限性

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

1. 无法修改存储决策：LSTM 的主要局限性之一是它难以修改存储的值，当遇到更相似的向量时。这可能导致在需要动态更新存储信息的任务中表现不佳。
2. 存储容量有限：LSTMs 将信息压缩到标量单元状态，这可能限制了它们有效存储和检索复杂数据模式的能力，特别是在处理罕见标记或长距离依赖时。
3. 缺乏 并行化：LSTMs 中的内存混合机制，涉及时间步之间的隐藏-隐藏连接，强制执行顺序处理，阻碍计算的并行化，限制了可扩展性。

这些局限性为变压器和其他架构的出现铺平了道路，这些架构在某些方面，尤其是在扩展到更大模型时，已经超越了 LSTMs。

xLSTM 架构

扩展 LSTM（xLSTM）家族

xLSTM 的核心是对传统 LSTM 框架的两个主要修改：指数门控和新颖的存储结构。这些增强引入了两种新的 LSTM 变体，分别称为 sLSTM（标量 LSTM）和 mLSTM（矩阵 LSTM）。

1. sLSTM：具有指数门控和内存混合的标量 LSTM
   • 指数门控：sLSTM 集成了输入和遗忘门的指数激活函数，实现对信息流的更灵活的控制。
   • 归一化和稳定：为了防止数值不稳定，sLSTM 引入了一个跟踪输入门和未来遗忘门乘积的归一化状态。
   • 内存混合：sLSTM 支持多个内存单元，并允许通过循环连接进行内存混合，实现复杂模式的提取和状态跟踪能力。
2. mLSTM：具有增强存储容量的矩阵 LSTM
   • 矩阵存储：mLSTM 使用矩阵存储代替标量存储，增加了存储容量，实现了信息的更高效检索。
   • 协方差更新规则：mLSTM 采用了协方差更新规则，受双向关联记忆（BAMs）的启发，高效地存储和检索键值对。
   • 并行化：通过放弃内存混合，mLSTM 实现了完全并行化，允许在现代硬件加速器（如 GPU）上进行高效计算，并实现了可扩展性。

这两个变体，sLSTM 和 mLSTM，可以集成到残差块架构中，形成 xLSTM 块。通过残差堆叠这些 xLSTM 块，研究人员可以构建强大的 xLSTM 架构，适用于特定的任务和应用领域。

数学

传统 LSTM：

原始的 LSTM 架构引入了恒定误差旋转木马和门控机制，以克服循环神经网络中的消失梯度问题。

LSTM 中的重复模块 – 来源

LSTM 内存单元更新由以下方程控制：

单元状态更新：ct = ft ⊙ ct-1 + it ⊙ zt

隐藏状态更新：ht = ot ⊙ tanh(ct)

其中：

• 𝑐𝑡​ 是时间 $t$ 的单元状态向量
• 𝑓𝑡​ 是遗忘门向量
• 𝑖𝑡 是输入门向量
• 𝑜𝑡 是输出门向量
• 𝑧𝑡 是由输入门调制的输入
• ⊙ 表示元素级乘法

门 𝑓𝑡、𝑖𝑡 和 𝑜𝑡 控制着哪些信息被存储、遗忘和从单元状态 𝑐𝑡 中输出，缓解了消失梯度问题。

具有指数门控的 xLSTM：

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) 或 ft = exp(W_f xt + R_f ht-1 + b_f)

输入（it）和遗忘（ft）门的指数激活函数，以及归一化状态 nt，使得对内存更新和修改存储信息的控制更加有效。

xLSTM 的关键特性和优势

1. 修改存储决策的能力：由于指数门控，xLSTM 可以有效地修改存储的值，当遇到更相关的信息时，克服了传统 LSTMs 的一个重大局限性。
2. 增强的存储容量：mLSTM 中的矩阵存储提供了增加的存储容量，使 xLSTM 能够更有效地处理罕见标记、长距离依赖和复杂的数据模式。
3. 并行化：mLSTM 变体的 xLSTM 是完全可并行的，允许在现代硬件加速器（如 GPU）上进行高效计算，并实现了可扩展性。
4. 内存混合和状态跟踪：sLSTM 变体的 xLSTM 保留了传统 LSTMs 的内存混合能力，实现了状态跟踪，使 xLSTM 比变压器和状态空间模型在某些任务中更具表现力。
5. 可扩展性：通过利用大型语言模型的最新技术，xLSTM 可以扩展到数十亿个参数，解锁语言建模和序列处理任务中的新可能性。

实验评估：展示 xLSTM 的能力

研究论文对 xLSTM 进行了全面实验评估，突出了其在各种任务和基准测试中的性能。以下是其中一些关键发现：

1. 合成任务和长距离竞技场：
   • xLSTM 在需要状态跟踪的形式语言任务中表现出色，超越了变压器、状态空间模型和其他循环神经网络架构。
   • 在多查询关联回忆任务中，xLSTM 展示了增强的存储容量，超越了非变压器模型，并且其性能与变压器相媲美。
   • 在长距离竞技场基准测试中，xLSTM 展示了稳定的强大性能，展示了其处理长距离问题的效率。
2. 语言建模和下游任务：
   • 当在 15B 个标记的 SlimPajama 数据集上进行训练时，xLSTM 在验证集上的困惑度方面超越了现有方法，包括变压器、状态空间模型和其他循环神经网络变体。
   • 随着模型规模的增加，xLSTM 继续保持其性能优势，展示了有利的扩展行为。
   • 在下游任务中，如常识推理和问答，xLSTM 在各种模型规模下表现为最佳方法，超越了最先进的方法。
3. 在 PALOMA 语言任务中的性能：
   • 在 PALOMA 语言基准测试的 571 个文本域中评估，xLSTM[1:0]（sLSTM 变体）在 99.5% 的域中实现了比 Mamba 更低的困惑度，在 85.1% 的域中实现了比 Llama 更低的困惑度，在 99.8% 的域中实现了比 RWKV-4 更低的困惑度。
4. 缩放法则和长度外推：
   • 当在 300B 个标记的 SlimPajama 数据集上进行训练时，xLSTM 展示了有利的缩放法则，表明其在模型规模增加时具有进一步提高性能的潜力。
   • 在序列长度外推实验中，xLSTM 模型即使在训练中看到的上下文明显更长的情况下，也保持了低困惑度，超越了其他方法。

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

实际应用和未来方向

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

1. 语言建模和文本生成：凭借其增强的存储容量和修改存储信息的能力，xLSTM 可能会革命性地改变语言建模和文本生成任务，使得文本生成更加连贯、上下文感知和流畅。
2. 机器翻译：xLSTM 的状态跟踪能力在机器翻译任务中可能非常有价值，在这些任务中，维护上下文信息和理解长距离依赖至关重要，以实现准确的翻译。
3. 语音识别和生成：xLSTM 的并行化和可扩展性使其非常适合语音识别和生成应用，在这些应用中，长序列的高效处理至关重要。
4. 时间序列分析和预测：xLSTM 处理长距离依赖和有效存储复杂模式的能力可能会在时间序列分析和预测任务中带来显著的改进，跨越金融、天气预报和工业应用等多个领域。
5. 强化学习和控制系统：xLSTM 在强化学习和控制系统中的潜力是有前途的，因为其增强的存储能力和状态跟踪能力可能会实现更智能的决策和在复杂环境中的控制。