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架构引入了恒定误差旋转和门控机制，以克服循环神经网络中的梯度消失问题。

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

实验评估：展示xLSTM的能力

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

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

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

实际应用和未来方向

xLSTM的潜在应用领域跨越了广泛的领域，从自然语言处理和生成到序列建模、时间序列分析和其他领域。以下是一些令人兴奋的领域，xLSTM可能会产生重大影响：

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