xLSTM: Extended Long Short-Term Memory — part 17/30

Cleaned-text code points [241404, 256449). Bibliography omitted here; source.md includes it.

VG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}{\hbox to47.78pt{\vbox to6.25pt{\pgfpicture\makeatletter\hbox{\hskip 23.88908pt\lower-1.94443pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-23.88908pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -33.06 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}{\hbox to33.45pt{\vbox to4.38pt{\pgfpicture\makeatletter\hbox{\hskip 16.72371pt\lower-1.36108pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-16.72371pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -23.14 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}{\hbox to27.26pt{\vbox to3.13pt{\pgfpicture\makeatletter\hbox{\hskip 13.6321pt\lower-0.9722pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-13.6321pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -18.86 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}\!\!\left(\tilde{\mathrm{f}}_{t}\right),\ \ \>\!\tilde{\mathrm{f}}_{t}\ =\ \bm{w}^{\top}_{\mathrm{f}}\ \bm{x}_{t}\ +\ b_{\mathrm{f}} forget gate (26) = σ ⁡ ( 𝐨 ~ t ) , 𝐨 ~ t = 𝑾 𝐨 ​ 𝒙 t + 𝒃 𝐨 \displaystyle=\ \sigma\left(\tilde{\mathbf{o}}_{t}\right)\ ,\qquad\qquad\quad\ \ \ \,\tilde{\mathbf{o}}_{t}\ =\ \bm{W}_{\mathbf{o}}\ \bm{x}_{t}\ +\ \bm{b}_{\mathbf{o}} output gate (27)

mLSTM can have multiple memory cells like the original LSTM. For mLSTM, multiple heads and multiple cells are equivalent as there is no memory mixing. In order to stabilize the exponential gates of mLSTM, we use the same stabilization techniques as for sLSTM (see Equation 15 ). Since the mLSTM has no memory mixing, this recurrence can be reformulated in a parallel version. For more details we refer to Appendix A.3 .

### 2.4 xLSTM Architecture

##### xLSTM Blocks.

An xLSTM block should non-linearly summarize the past in a high-dimensional space to better separate different histories or contexts. Separating histories is the prerequisite to correctly predict the next sequence element such as the next token. We resort to Cover’s Theorem ( Cover, 1965 ) , which states that in a higher dimensional space non-linearly embedded patterns can more likely be linearly separated than in the original space. We consider two residual block architectures: (i) A residual block with post up-projection (like Transformers), which non-linearly summarizes the past in the original space, then linearly maps into a high-dimensional space, applies a non-linear activation function, and linearly maps back to the original space; see left panel of Figure 3 and third column in Figure 1 . A more detailed version is depicted in Figure 10 in the appendix. (ii) A residual block with pre up-projection (like State Space Models), which linearly maps to a high-dimensional space, non-linearly summarizes the past in the high-dimensional space and then linearly maps back to the original space. For an xLSTM block containing an sLSTM, we use the post up-projection block. For an xLSTM block containing an mLSTM, we use the pre up-projection block since the memory capacity becomes larger in the high-dimensional space. We refer to the left panel of Figure 3 and third column in Figure 1 , or Figure 10 in the appendix for more details.

##### xLSTM Architecture.

An xLSTM architecture is constructed by residually stacking building blocks ( Srivastava et al., 2015 ; He et al., 2016 ) . We rely on the most commonly used pre-LayerNorm ( Ba et al., 2016b ) residual backbones as used in contemporary Large Language Models. See last column in Figure 1 .

### 2.5 Memory and Speed Considerations

Contrary to Transformers, xLSTM networks have a linear computation and a constant memory complexity with respect to the sequence length. Since the xLSTM memory is compressive, it is well suited for industrial applications and implementations on the edge.

The memory of mLSTM does not require parameters but is computationally expensive through its d × d d\times d matrix memory and d × d d\times d update. We trade off memory capacity against computational complexity. Nevertheless, the computations can be done in parallel on GPUs, therefore these computations have only a minor effect on the wall clock time.

While mLSTM is parallelizable analog to FlashAttention ( Dao et al., 2022 ; Dao, 2024 ) or GLA ( Yang et al., 2023 ) , sLSTM is not parallelizable due to the memory mixing (hidden-hidden connections). However, we developed a fast CUDA implementation with GPU memory optimizations to the register level which is typically less than two times slower than mLSTM.

## 3 Related Work

##### Linear Attention.

Several methods have been suggested to overcome the quadratic complexity in terms of context length of the Transformer and make attention linear in the context length. The Synthesizer learns synthetic attention weights without token–token interactions ( Tay et al., 2020 ) . Linformer realizes self-attention by a low-rank matrix and even linearly approximates it ( Wang et al., 2020 ) . Linear Transformer linearizes the attention mechanism ( Katharopoulos et al., 2020 ) . Performer linearly approximates the attention softmax by positive orthogonal random features approach ( Choromanski et al., 2021 ) . Attention has been replaced by fast long convolutions in the Structured Global Convolution (SGConv) ( Li et al., 2022 ) and the Hyena Hierarchy ( Poli et al., 2023 ) .

##### State Space Models.

Recently, State Space Models (SSMs) became very popular since they are linear in the context length and show promising performance compared to Transformers. One of the first proposed models was Structured State Space sequence model (S4) ( Gu et al., 2021 ) , followed by Diagonal State Space (DSS) model ( Gupta et al., 2022 ) , Gated State Space (GSS) models ( Mehta et al., 2022 ) , S5 model ( Smith et al., 2022 ) , Bidirectional Gated SSM (BiGS) ( Wang et al., 2022 ) , H3 model ( Fu et al., 2023 ) , and Mamba ( Gu & Dao, 2023 ) .

##### Recurrent Neural Networks.

Recurrent Neural Networks (RNNs) have been suggested to replace Transformer and attention due to their linearity in the context length. RNNs with Deep Linear Recurrent Units (LRUs) showed promising results for language modeling ( Orvieto et al., 2023 ; De et al., 2024 ) , as did Hierarchically Gated Linear RNN (HGRN) ( Qin et al., 2023 ) and HGRN2 ( Qin et al., 2024 ) . A well-known RNN approach to large language modeling is RWKV ( Peng et al., 2023 ; Peng et al., 2024 ) , showcasing competitive performance to Transformers.

##### Gating.

One of the key ideas of LSTM is gating, which was rediscovered and reinterpreted in many recent approaches. Gating was used in HGRN ( Qin et al., 2023 ) , HGRN2 ( Qin et al., 2024 ) , Gated Linear Attention (GLA) ( Yang et al., 2023 ) , Gated State Space (GSS) models ( Mehta et al., 2022 ) , Bidirectional Gated SSM (BiGS) ( Wang et al., 2022 ) , Moving Average Equipped Gated Attention (MEGA) ( Ma et al., 2022 ) , RWKV ( Peng et al., 2023 ) , and Mamba ( Gu & Dao, 2023 ) .

##### Covariance Update Rule.

To enhance storage capacities, we equipped the mLSTM cell with a matrix memory with a covariance update rule. Other methods which build on such an update mechanism are Fast Weight Programmers ( Schmidhuber, 1992 ; Schlag et al., 2021 ) , RWKV-5 and RWKV-6 ( Peng et al., 2024 ) , Retention ( Sun et al., 2023 ) , Linear Transformer ( Katharopoulos et al., 2020 ) , and HGRN2 ( Qin et al., 2024 ) .

##### Most Related.

Conceptually the closest models to xLSTM are Retention ( Sun et al., 2023 ) , RWKV ( Peng et al., 2023 ; Peng et al., 2024 ) , and HGRN2 ( Qin et al., 2024 ) . These models share the concepts matrix memory and/or gating. However, in contrast to the new sLSTM, these approaches do not allow memory mixing. Memory mixing enables to solve state tracking problems, and therefore LSTMs are more expressive than State Space Models (SSMs) and Transformers ( Merrill et al., 2024 ; Delétang et al., 2023 ) . State tracking is required to evaluate code or to track entities in a long narrative.

##### Residually Stacking Architectures.

Like almost all contemporary large deep learning models, xLSTM architectures are constructed by residually stacking building blocks ( Srivastava et al., 2015 ; He et al., 2016 ) . This construction enabled deep convolutional networks ( He et al., 2016 ) and Transformers ( Vaswani et al., 2017 ) . Transformers are the ultimate force behind Large Language Models (LLMs) like GPT-3 ( Brown et al., 2020 ) , ChatGPT ( Schulman et al., 2022 ) , GPT-4 ( Achiam et al., 2023 ) , Megatron-LM ( Shoeybi et al., 2019 ) , Gopher ( Rae et al., 2021 ) , ERNIE 3.0 Titan ( Wang et al., 2021 ) , GLaM ( Du et al., 2021 ) , Chinese M6 ( Lin et al., 2021 ) , mutilingual AlexaTM 20B ( Soltan et al., 2022 ) , OPT ( Zhang et al., 2022 ) , Chinchilla ( Hoffmann et al., 2022 ) , BLOOM ( Scao et al., 2022 ) , GLM-130B ( Zeng et al., 2022 ) , LaMDA ( Thoppilan et al., 2022 ) , PaLM ( Chowdhery et al., 2022 ) , Llama ( Touvron et al., 2023 ) , Gemini ( Google, 2023 ; Reid et al., 2024 ) .

## 4 Experiments

In this section, we experimentally evaluate xLSTM and compare it to existing methods with a focus on language modeling. We investigate xLSTM’s specific capabilities on synthetic tasks in Section 4.1 . In Section 4.2 , we compare the validation set perplexity of various current language modeling methods that were trained on 15B tokens from SlimPajama ( Soboleva et al., 2023 ) . On the same dataset, we perform ablation studies for xLSTM. Then, we assess the scaling behavior of the different methods analogous to Kaplan et al. (2020) and Brown et al. (2020) . In Section 4.3 , we conduct a more thorough language modeling experiment. We compare xLSTM and the best performing methods from Section 4.2 after being trained on 300B tokens from SlimPajama ( Soboleva et al., 2023 ) . First, we assess how well the methods perform in extrapolating to longer contexts, secondly we test the methods via validation perplexity and performance on downstream tasks ( Sutawika et al., 2024 ) , thirdly we evaluate the methods on 571 text domains of the PALOMA language benchmark dataset ( Magnusson et al., 2023 ) , fourthly we again assess the scaling behavior of the different methods, but now with 20 times more training data.

For all experiments, we use the notation xLSTM[ a a : b b ] for the ratio a / b a/b of mLSTM-based versus sLSTM-based xLSTM blocks. For example, xLSTM[7:1] means that out of eight blocks, seven are mLSTM-based blocks and one is an sLSTM-based block. For a common total block number of 48, this translates to 6 sLSTM-based blocks and 42 mLSTM-based blocks. Further, for all experiments, we use pre and post up-projection blocks for mLSTM and sLSTM, respectively.

### 4.1 Synthetic Tasks and Long Range Arena

First, we test the effectiveness of xLSTM’s new exponential gating with memory mixing on formal languages ( Delétang et al., 2023 ) . Then, we assess the effectiveness of xLSTM’s new matrix memory on the Multi-Query Associative Recall task ( Arora et al., 2023 ) . Finally, xLSTM’s performance at processing long sequences in the Long Range Arena is evaluated ( Tay et al., 2021 ) .

##### Test of xLSTM’s Exponential Gating with Memory Mixing.

We test xLSTM’s new exponential gating with memory mixing, which should enable it to solve state tracking problems ( Merrill et al., 2024 ; Merrill & Sabharwal, 2023 ) . We implement and extend the formal language tasks from Delétang et al. (2023) to enable multi-length training for length extrapolation. For a detailed description of all tasks and extended results see Appendix B.1.1 . We compare xLSTM to other methods including Transformers, State Space Models, and Recurrent Neural Networks. The accuracy of the tested methods is evaluated on those tokens relevant to the task. The accuracy is scaled between 0 (random) and 1 (perfect). We compare 2-block architectures of the following methods on these tasks: xLSTM[0:1] (i.e., only sLSTM), xLSTM[1:0] (i.e., only mLSTM), xLSTM[1:1], Llama, Mamba, RWKV, Retention, Hyena, LSTM, and LSTM in Transformer blocks (LSTM (Block)). The results of this experiment are shown in Figure 4 . Models such as Transformers or State Space Models without memory mixing (no state tracking) cannot solve, e.g. regular grammars like the parity task. This result is in agreement with findings that Transformers and State Space models are fundamentally less powerful than RNNs ( Merrill et al., 2024 ; Merrill & Sabharwal, 2023 ; Delétang et al., 2023 ) .

##### Test of xLSTM’s Memory Capacities on Associative Recall Tasks.