xLSTM: Extended Long Short-Term Memory

引用： Beck M, Pöppel K, Spanring M, et al. xLSTM: Extended Long Short-Term Memory[J]. arXiv preprint arXiv:2405.04517, 2024.

论文链接： [2405.04517] xLSTM: Extended Long Short-Term Memory (arxiv.org)

作者： Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, Sepp Hochreiter

机构： ELLIS Unit, LIT AI Lab, Institute for Machine Learning, JKU Linz, Austria; NXAI Lab, Linz, Austria; NXAI GmbH, Linz, Austria

摘要

• 论文提出了xLSTM，一种扩展的长短期记忆网络，旨在解决传统LSTM的局限性，并在大规模参数下进行语言建模。
• xLSTM引入了指数门控和适当的归一化与稳定技术，修改了LSTM记忆结构，包括标量记忆的sLSTM和完全可并行化的具有矩阵记忆和协方差更新规则的mLSTM。
• 通过将这些LSTM变体集成到残差块中，构建了xLSTM架构，这些架构在性能和扩展性方面与最先进的Transformers和状态空间模型相媲美。

引言

• LSTM自1990年代引入以来，在多个领域取得了成功，特别是在大型语言模型（LLMs）中。
• 引入Transformer技术后，其并行化的自注意力机制使得LSTM在大规模应用中的性能受到挑战。
• 论文提出了一个问题：当LSTM扩展到数十亿参数，并结合现代LLMs的最新技术，同时克服LSTM的已知限制时，我们能在语言建模中走多远？

xLSTM架构

1. sLSTM（Scalar LSTM）

指数门控是sLSTM中的一个创新点，它允许模型更有效地更新其记忆状态。在传统的LSTM中，门控机制通常涉及sigmoid函数，但在xLSTM中，输入门（$i_t$）和遗忘门（$f_t$）可以具有指数激活函数:

$$c_t=f_t{c}_{t-1}+i_t{z}_t$$

$$n_t=f_t{n}_{t-1}+i_t$$

$$h_t=o_t\widetilde{h_t},\widetilde{h_t}=o_t/n_t$$

$$z_t=\phi ({\tilde{z}}_t),{\tilde{z}}_t=w_z^T{x}_t+r_z{h}_{t-1}+b_z$$

$$i_t=exp({\tilde{i}}_t),{\tilde{i}}_t=w_i^T{x}_t+r_i{h}_{t-1}+b_i$$

$$f_t=\sigma ({\tilde{f}}_t)ORexp({\tilde{f}}_t),{\tilde{f}}_t=w_f^T{x}_t+r_f{h}_{t-1}+b_f$$

$$o_t=exp({\tilde{o}}_t),{\tilde{o}}_t=w_o^T{x}_t+r_o{h}_{t-1}+b_o$$

指数激活函数可能导致较大的值，从而导致溢出。因此，用一个额外的状态$m_t$来稳定门：

$$m_t=\max (\log (f_t)+m_{t-1},\log (i_t))$$

$$i_t^{\prime }=exp(\log (i_t)-m_t)=exp({\tilde{i}}_t-m_t)$$

$$f_t^{\prime }=exp(\log (f_t)+m_{t-1}-m_t)$$

其中，$m_t​ $是稳定状态，用于防止梯度爆炸。

同时，sLSTM引入了新的记忆混合技术，允许在多个内存单元之间进行更复杂的交互。多个存储器单元使得能够分别经由从隐藏状态向量$h$到存储器单元输入$z$和门$i$、$f$、$o$的循环连接$R_z$、$R_i$、$R_f$、$R_o$进行存储器混合。sLSTM可以有多个头，每个头内混合内存，但不能跨头混合。

2. mLSTM（Matrix LSTM）

mLSTM使用矩阵记忆来增强存储容量，并通过协方差更新规则来存储关键值对。

$$C_t=f_t{C}_{t-1}+i_t{v}_t{k}_t^T$$

$$n_t=f_t{n}_{t-1}+i_t{k}_t$$

$$h_t=o_t\odot {\tilde{h}}_t,{\tilde{h}}_t=C_t{q}_t/\max \left\{|n_t^T{q}_t|,1\right\}$$

$$q_t=W_q{x}_t+b_q$$

$$k_t=\frac{1}{\sqrt{d}}{W}_k{x}_t+b_k$$

$$v_t=W_v{x}_t+b_v$$

$$i_t=exp({\tilde{i}}_t),{\tilde{i}}_t=w_i^T{x}_t+b_i$$

$$f_t=\sigma ({\tilde{f}}_t)ORexp({\tilde{f}}_t),{\tilde{f}}_t=w_f^T{x}_t+b_f$$

$$o_t=\sigma ({\tilde{o}}_t),{\tilde{o}}_t=w_o^T{x}_t+b_o$$

3. xLSTM块（xLSTM Blocks）

xLSTM块结合了sLSTM和mLSTM的特性，并通过残差连接来进一步提高性能。对于残差sLSTM块，输入首先进入sLSTM，然后是一个门控的多层感知机（MLP）。对于残差mLSTM块，输入首先通过两个MLP，然后是mLSTM，通过卷积、可学习的跳跃连接和输出门。

4. xLSTM架构（xLSTM Architecture）

xLSTM架构通过残差堆叠xLSTM块来构建，利用了预层归一化（preLayerNorm）残差骨干网络。

实验

• 论文在合成任务和长距离竞技场（Long Range Arena）上测试了xLSTM，并与其他方法进行了比较。
• 在SlimPajama数据集上进行了语言建模实验，比较了不同方法的性能。
• 进行了扩展实验，训练了更大的模型，并在更多的训练数据上评估了它们的扩展行为。

结论

• xLSTM在语言建模方面的表现至少与当前的Transformer或状态空间模型相当。
• xLSTM有潜力在强化学习、时间序列预测或物理系统建模等深度学习领域产生重大影响。