LSTM原作者携改进版XLSTM归来以On复杂度挑战Transformers的O

🔔 LSTM 的 发明者（Sepp Hochreiter）推 出了一种新的 LLM 架构！—— XLSTM

其主要组成部分是一种新的可并行化 LSTM。

⚠️注：以前的 LSTM 的主要弱点是它们的序列性（无法同时进行），LSTM 的三个主要弱点使得Transformers 更为优秀：

无法修改存储决策

存储容量有限

由于内存混合，缺乏并行性

如果你不明白，那很正常不是你的问题是我的问题，请看 GIF 展示：

gif 来源：https://towardsdatascience.com/illustrated-guide-to-lstms-and-gru-s-a-step-by-step-explanation-44e9eb85bf21

可以看到 LSTM 是顺序的，这意味着必须一个接一个地经过绿色盒子（这里有所简化）。在可以继续之前，需要拿到前一个盒子的结果。

而谷歌论文《Attention Is All You Need》中的Transformers 不这样做。它们在标记之间并行化操作，这是一个非常伟大的事情。

Sepp 和团队的新解决方案？

论文链接：https://arxiv.org/pdf/2405.04517

主要贡献：

LSTM 的矩阵内存（无内存混合）

指数门控

这可能对你来说没什么意义，且听分解！👇

他们对输入门和遗忘门使用指数激活，加上一个正则化状态。这有助于 LSTM 在看到新数据时更有效地修订其内存。

左侧：原始 LSTM 更新规则

右侧：新的 sLSTM 更新规则

接下来，他们用矩阵内存取代了标量内存单元。

在常规的 LSTM 中，内存单元是单个标量值。这限制了 LSTM 在每个时间步可以存储的信息量。xLSTM 用一个完整的矩阵替换了这个标量，可以捕获更多信息。

代替标量单元状态 c_t，我们现在有一个大小为 d×d 的矩阵 C_t，其中 d 是隐藏状态大小。在每个时间步，我们希望存储一个键向量 k_t 和一个值向量 v_t，两者的大小都是 d。矩阵内存的更新规则是：C_t = f_t * C_{t-1} + i_t * (v_t * k_t^T)

让我们分解一下：f_t 是遗忘门，控制保留多少以前的记忆。i_t 是输入门，控制添加多少新信息。v_t * k_t^T 是键和值向量的外积，结果是一个 d×d 矩阵。

这种外积操作是矩阵内存工作的关键。它允许 LSTM 存储和检索大量的键值对！

当我们想从矩阵内存中检索一个值时，我们使用当前的隐藏状态 h_t 作为查询向量 q_t。检索的值是：h_t = o_t * (C_t * q_t) / max(norm(q_t), 1)

这里，o_t 是输出门，控制内存的输出，而 max 操作是为了稳定性。直观地说，矩阵内存允许 LSTM 软性地关注 C_t 中存储的大量键值对集，就像变压器的注意力机制工作方式一样。查询向量 q_t 与 C_t 中所有存储的键进行比较，相应的值根据相似性被检索并组合。

这大大增加了 LSTM 的存储容量，因为它现在可以在每个时间步存储多达 d 个不同的键值对，而不仅仅是单个标量。这使得 xLSTM 能够从输入序列中记住更多信息，这对于语言建模特别有用，因为存在许多长期依赖性。

LSTM 的清晰解释由 statQuest 提供：https://youtube.com/watch?v=YCzL96nL7j0

O(N) 时间复杂度 vs Transformers  O(N^2) 复杂度

重要的是，作者指出 xLSTM 的时间复杂度为 O(N)，空间复杂度为 O(1)，随着序列长度的增加。变压器的时间和空间复杂度是二次的 O(N^2)，所以 xLSTM 更好。

尽管如此，请注意，已经有很多改进，使注意力更高效，比如 HYENA 和 Sparse Attention，这可以将复杂度降低到更像 O(NlogN) 或 O(Nsqrt(N))

那么，xLSTM 的表现如何？

作者团队进行了许多复杂的评估，但我觉得这个最清晰。他们训练了几个模型（变压器 LLM、RWKV 和 XLSTM）在 150 亿个文本标记上，并测量了困惑度（每个模型预测下一个标记的好坏）

xLSTM[1:0] (1 mLSTM, 0 sLSTM 块) 表现最好

此外，XLSTM 遵循与传统基于Transformers LLM 相似的规模化法则。