【论文阅读】xLSTM: Extended Long Short-Term Memory

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^{}$$