程序员学长 | 快速学会一个算法，xLSTM

本文来源公众号“程序员学长”，仅用于学术分享，侵权删，干货满满。

原文链接：快速学会一个算法，xLSTM

今天给大家分享一个超强的算法模型，xLSTM。

xLSTM（Extended Long Short-Term Memory）是对传统 LSTM（Long Short-Term Memory）模型的扩展和改进，旨在提升其在处理时间序列数据和序列预测任务中的性能。

传统的 LSTM 模型通过引入遗忘门、输入门和输出门，解决了标准 RNN（Recurrent Neural Network）在长序列数据中存在的梯度消失和梯度爆炸问题。然而，LSTM 仍然存在一些局限性，尤其是在处理非常长的序列或需要更高维度特征提取时。xLSTM 的提出是为了克服这些限制。

xLSTM 中的核心增强功能

具体改进包括如下几个方面。

指数门控

xLSTM 论文中引入的指数门控机制是对 LSTM 中使用的传统 S 形门控的重大改进。通过对输入门和遗忘门采用指数激活，xLSTM 增强了模型在处理新信息时有效修改和更新其记忆的能力。

在传统的 LSTM 中，S 型门控函数限制了模型对记忆单元状态进行重大改变的能力，尤其是当函数值接近 0 或 1 时。这种限制阻碍了 LSTM 快速适应新数据的能力，并可能导致内存更新不理想。

xLSTM 通过将 S 型激活替换为指数激活来解决此问题。指数门控允许记忆细胞状态发生更明显的变化，使模型能够快速整合新信息并相应地调整其记忆。规范化器状态有助于稳定指数门控并保持输入门和遗忘门之间的平衡。
先进的内存结构

xLSTM 采用了 sLSTM 和 mLSTM 等先进的内存设计，它们分别使用了新的混合技术和基于矩阵的内存。这些结构允许更好的并行处理能力和更高效的数据存储和检索，这对于处理大规模数据至关重要。
残差连接

通过在 LSTM 单元内集成残差块，xLSTM 可以维持甚至增强整个网络的梯度流，这对于有效训练更深层的模型至关重要。
可并行架构

xLSTM 论文中最重要的进步之一是引入了可并行化的架构，这解决了传统 LSTM 的一个主要限制。

在传统的 LSTM 中，标记的处理是按顺序进行的，每个标记一次处理一个，这限制了模型利用并行性的能力，并导致训练和推理时间变慢。

xLSTM 架构引入了 mLSTM（矩阵内存 LSTM）和 sLSTM（标量 LSTM）块的灵活组合，从而实现了对 token 的并行处理。mLSTM 块旨在同时对整个 token 序列进行操作，从而实现与 Transformer 模型实现的并行性类似的高效并行计算。

mLSTM 模块利用了前面讨论过的矩阵内存机制，使其能够并行捕获和处理所有 token 中丰富的高维信息。这种并行处理能力显著加快了训练和推理过程，使得 xLSTM 比传统 LSTM 具有更高的计算效率。

另一方面，sLSTM 块保留了传统 LSTM 的顺序处理特性，使模型能够捕获对特定任务可能很重要的某些顺序依赖关系。

在 xLSTM 架构中，可以灵活地以不同的比例组合 mLSTM 和 sLSTM 块，从而在并行性和顺序建模之间取得平衡，从而能够适应各种语言建模任务。

关于 xLSTM 更详细的说明，请参考对应的论文

https://arxiv.org/pdf/2405.04517

LSTM 基础知识

为了解释 xLSTM，我们首先简单回顾一下 LSTM。

原始 LSTM 的计算涉及几个关键组件：输入门、遗忘门、输出门和单元状态。

以下是 LSTM 单元的标准计算步骤。

xLSTM 解释

之所以被称为 xLSTM，是因为它将原始 LSTM 扩展为 LSTM 的多种变体，例如 sLSTM 和 mLSTM，每种变体都针对特定的性能和功能进行了优化，以处理各种复杂的序列数据问题。

1.sLSTM

sLSTM 在 LSTM 的基础上添加了标量更新机制。该设计通过对内部存储单元进行细粒度控制来优化门控机制，使其更适合处理具有细微时间变化的序列。

sLSTM 通常利用指数门控和归一化技术来提高处理长序列数据的稳定性和准确性。通过这种方法，sLSTM 能够提供与更复杂模型相当的性能，同时保持较低的计算复杂度，使其特别适合资源受限的环境或需要快速响应的应用程序。

2.mLSTM

mLSTM（矩阵 LSTM）通过将原始 LSTM 中的向量运算扩展为矩阵运算，显著增强了模型的内存容量和并行处理能力。

在 mLSTM 中，每个状态不再是单个向量，而是一个矩阵。这使得它能够在单个时间步内捕获更复杂的数据关系和模式。mLSTM 特别适合处理大规模数据集或需要识别高度复杂数据模式的任务。

将矩阵引入状态表示使得 mLSTM 能够更自然、更有效地处理多维数据。通过以矩阵形式处理数据，mLSTM 可以同时处理多个数据点，从而提高吞吐量和学习过程的速度。这种能力在图像和视频处理等领域尤其有价值，因为这些领域的数据本质上以矩阵形式存在。

效率与性能分析

xLSTM 论文对所提架构的效率和性能进行了全面分析，突出了其相对于基于 Transformer 的模型的优势。作者进行了一系列实验和比较，以证明 xLSTM 卓越的计算效率和建模能力。

xLSTM 的一个关键效率优势在于其时间和内存复杂度。传统的基于 Transformer 的模型相对于序列长度 N 表现出 O(N²) 的二次方时间和内存复杂度。这意味着随着序列长度的增加，Transformer 的计算成本和内存需求呈二次方增长，使其处理长序列的效率较低。

相比之下，xLSTM 实现了线性时间复杂度 O(N) 和恒定内存复杂度 O(1)（相对于序列长度）。这比 Transformers 有了显著的改进，因为它允许 xLSTM 更有效地处理更长的序列，而无需二次方增加计算成本和内存使用量。线性时间复杂度可以缩短训练和推理时间，而恒定内存复杂度确保即使对于长序列，内存需求仍然是可控的。

为了验证效率和性能声明，作者通过在包含 150 亿个 token 的大规模数据集上训练多个模型进行了比较评估。评估中包括的模型是基于 Transformer 的语言模型 (LLM)、RWKV 模型和 xLSTM 的不同变体。

评估结果为 xLSTM 的卓越性能提供了有力证据。特别是，由一个 mLSTM 块和零个 sLSTM 块组成的 xLSTM[1:0] 变体在所有测试模型中实现了最低的困惑度。困惑度是语言建模中广泛使用的指标，用于衡量模型预测序列中下一个标记的能力。困惑度越低，语言建模性能越好。