LSTM一作新作xLSTM架构：大幅领先Transformer和状态空间模型（SSM）

这篇论文介绍了一种名为xLSTM（Extended Long Short-Term Memory）的新型递归神经网络架构，旨在解决传统LSTM（Long Short-Term Memory）网络的一些局限性，并提高其在语言建模等任务中的性能。

论文：xLSTM: Extended Long Short-Term Memory 链接：https://arxiv.org/pdf/2405.04517

下面是对论文各部分的详细解读。

摘要

• 背景：LSTM在1990年代被提出，用以解决循环神经网络（RNN）的梯度消失问题。LSTM在多种领域取得了成功，但随着Transformer技术的出现，其地位受到了挑战。
• 问题：作者提出了一个问题：如果将LSTM扩展到数十亿参数，并利用现代大型语言模型（LLM）的技术，同时克服LSTM的已知限制，我们能在语言建模上走多远？

• 贡献：论文介绍了两种新的LSTM变体：sLSTM（具有标量记忆和更新）和mLSTM（具有矩阵记忆和协方差更新规则），并将它们集成到残差块中，形成xLSTM架构。

引言

• LSTM原理：介绍了LSTM的核心原理，包括恒定误差旋转（constant error carousel）和门控机制。
• LSTM应用：LSTM在文本生成、序列到序列翻译、程序评估等任务中表现出色。
• LSTM局限性：指出LSTM的三个主要局限性：无法修订存储决策、有限的存储容量、缺乏可并行化。

扩展的长短期记忆

• sLSTM：引入了指数门控和新的存储混合技术，允许LSTM修订其存储决策。
• mLSTM：将LSTM的记忆单元从标量扩展到矩阵，提高了存储容量，并引入了协方差更新规则，使得mLSTM可以完全并行化。
• xLSTM架构：通过将sLSTM和mLSTM集成到残差块中，构建了xLSTM架构。

记忆单元

• 恒定误差旋转：ct = ft * ct-1 + it * zt，其中ct是单元状态，ft是遗忘门，it是输入门，zt是经过激活函数的输入。
• 门控：包括输入门it、遗忘门ft和输出门ot，控制信息的流动。

sLSTM

• 指数门控：引入指数激活函数到输入和遗忘门中。
• 标准化和稳定化：通过规范化器状态来平衡门控的激活值。

mLSTM

• 矩阵记忆：使用矩阵来存储和检索信息，提高了存储容量。
• 协方差更新规则：使用协方差矩阵来更新记忆单元。

xLSTM架构

• 残差块：xLSTM架构通过残差堆叠构建，利用预层归一化（preLayerNorm）残差骨干。
• 并行化：mLSTM的设计允许并行化，而sLSTM由于内存混合而无法并行化。

实验

• 合成任务和长距离竞技场：测试了xLSTM在处理形式语言和长序列任务上的能力。

• 方法比较和消融研究：在SlimPajama数据集上训练并比较了xLSTM和其他方法的性能。

• 大型语言模型：在更大的数据集上训练xLSTM，并评估其作为大型语言模型的潜力。

相关工作

• 线性注意力：讨论了几种旨在降低Transformer注意力机制复杂度的方法。
• 状态空间模型：最近在语言建模中变得流行的方法，它们在上下文长度上是线性的。

结论

• 性能：xLSTM在语言建模任务上与现有的Transformer和状态空间模型相比表现出色。
• 潜力：xLSTM有潜力成为强化学习、时间序列预测或物理系统建模等领域的重要工具。

限制

• 并行化：sLSTM的内存混合阻止了并行化操作。
• 计算复杂性：mLSTM的矩阵记忆具有高计算复杂性。
• 初始化：遗忘门的初始化需要谨慎选择。
• 内存限制：矩阵记忆可能在长序列上下文中超载。

这篇论文的每个部分都详细介绍了xLSTM架构的设计原理、数学公式和实验结果，展示了其在语言建模任务中的潜力。