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

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

In this experiment, we test xLSTM’s new matrix memory in terms of the memory capacity on the Multi-Query Associative Recall task ( Arora et al., 2023 ) : For each sequence, key-value pairs are randomly chosen from a large vocabulary, which must be memorized for later retrieval. To enhance the difficulty of the original task, we increase the number of key-value pairs up to 256 and extend the context length up to 2048. Thus, we have broader tests for the memory capacities of different models. We compare 2-block architectures of Llama, Mamba, RWKV-5, RWKV-6, xLSTM[1:1] and xLSTM[1:0]. The models are evaluated by the accuracy at recalling the pairs. Since Transformers (e.g. Llama) have a memory that is exponential in the coding dimension ( Ramsauer et al., 2021 ) , they constitute the gold standard at this task. Results are shown in Figure 5 . xLSTM[1:1] performs best among all non-Transformer models, also for small models. Interestingly, the sLSTM block does not diminish the memory capacity but rather leverages it, which becomes evident at the most difficult task with 256 key-value pairs. Additional results are presented in Appendix B.1.2 , where extrapolation analyses indicate that xLSTM’s enhanced memory capacities also allow for extrapolating to contexts that are longer than those seen during training.

##### Test of xLSTM’s Long Context Capabilities on Long Range Arena.

To assess xLSTM’s performance on long sequences and large contexts, we compare different methods on the Long Range Arena ( Tay et al., 2021 ) . xLSTM demonstrates consistent strong performance on all of the tasks, suggesting that the xLSTM architecture is remarkably efficient in handling different aspects of long context problems. For more details, see Appendix B.1.3 .

### 4.2 Method Comparison and Ablation Study

To address the main question of our paper, i.e. what can our new LSTM variants achieve when scaled up in language modelling, we train xLSTMs, Transformers, State Space Models, and other methods on 15B tokens from SlimPajama in the same auto-regressive setting. We compare the trained models on the validation set and perform ablation studies for the xLSTMs.

##### Comparing xLSTM to Other Methods.

For comparison, we train models on 15B tokens from SlimPajama ( Soboleva et al., 2023 ) . The trained models are evaluated by their perplexity on the validation set. We compare the following methods: xLSTM (our new method), GPT-3 (Transformer) ( Brown et al., 2020 ) , Llama (Transformer) ( Touvron et al., 2023 ) , H3 (SSM) ( Fu et al., 2023 ) , Mamba (SSM) ( Gu & Dao, 2023 ) , RWKV-4 (RNN) ( Peng et al., 2023 ) , RWKV-5 (RNN) ( Peng et al., 2024 ) , RWKV-6 (RNN) ( Peng et al., 2024 ) , GLA (linear Transformer) ( Yang et al., 2023 ) , HGRN (RNN) ( Qin et al., 2023 ) , HGRN2 (RNN) ( Qin et al., 2024 ) . RetNet (linear Transformer) ( Sun et al., 2023 ) , Hyena (linear Transformer) ( Poli et al., 2023 ) , xLSTM[1:0], and xLSTM[7:1] (see Section 4 ). The models were trained with mixed precision, for RWKV-5, RWKV-6, GLA, HGRN2, the mixed- precision training did not utilize the PyTorch automated mixed precision (see also Appendix Section B.2 ). We categorize the methods into (a) Transformers, (b) State Space Models (SSMs), and (c) Recurrent Neural Networks (RNNs) together with linear Transformers. Linear Transformers are linear methods that substitute the Transformer attention mechanism. The models match a GPT-3 model with 350M parameters in size, i.e. embedding dim 1024 and 24 residual blocks. Only GPT-3 uses shared weights for token and output embeddings, therefore has fewer parameters. The results in Table 1 show that xLSTM outperforms all existing methods in validation perplexity. For details see Appendix B.2 . Figure 6 shows the scaling behaviour for this experiment, indicating that xLSTM will also perform favorably for larger models.

##### Ablation Studies.

Table 1 and Figure 6 demonstrate that xLSTM achieves excellent results at language modeling when being trained on 15B tokens from SlimPajama. To ablate the changes from LSTM to xLSTM, we morph a vanilla LSTM architecture step-by-step into an xLSTM architecture. Firstly, we integrate LSTM layers into pre-LayerNorm residual backbones. Secondly, we extend this to a post up-projection block. Finally, we add exponential gating and matrix memory. The results are shown in Table 2 (top). The ablation studies attribute the strong performance improvement to both the exponential gating and the matrix memory. Additionally, due to the importance of gating in RNNs and State Space Models, we ablate different gating mechanisms. In Table 2 (bottom), we conclude that having each gate learnable and influenced by the input has an incremental positive effect. Additional studies on the individual backbone components are discussed in Appendix B.2 .

Model Modification Exponential Gating Matrix Memory #Params M SlimPajama (15B) ppl ↓ \downarrow LSTM Vanilla Multi-Layer LSTM ✗ ✗ 607.8 2417.86 Adding Resnet Backbone ✗ ✗ 506.1 35.46 Adding Up-Projection Backbone ✗ ✗ 505.9 26.01 xLSTM[0:1] Adding Exponential Gating ✓ ✗ 427.3 17.70 xLSTM[7:1] Adding Matrix Memory ✓ ✓ 408.4 13.48

Learnable Gates Forget Gate Input Gate SlimPajama (15B) ppl ↓ \downarrow Input Dependent Learnable Bias Bias Init Input Dependent Learnable Bias Bias Init No Gates ✗ ✗ + ∞ +\infty ✗ ✗ 0 0 NaN No Gates ✗ ✗ [ 3 , 6 ] [3,6] ✗ ✗ 0 0 13.95 Forget Gate ✓ ✓ [ 3 , 6 ] [3,6] ✗ ✗ 0 0 13.58 Input Gate ✗ ✗ [ 3 , 6 ] [3,6] ✓ ✓ 𝒩 ⁡ ( 0 , 0.1 ) \mathcal{N}(0,0.1) 13.69 Forget Gate Bias ✗ ✓ [ 3 , 6 ] [3,6] ✗ ✗ 0 0 13.76 Forget + Input Gate Bias ✗ ✓ [ 3 , 6 ] [3,6] ✗ ✓ 𝒩 ⁡ ( 0 , 0.1 ) \mathcal{N}(0,0.1) 13.73 Forget Gate + Input Gate Bias ✓ ✓ [ 3 , 6 ] [3,6] ✗ ✓ 𝒩 ⁡ ( 0 , 0.1 ) \mathcal{N}(0,0.1) 13.55 Forget Gate + Input Gate ✓ ✓ [ 3 , 6 ] [3,6] ✓ ✓ 𝒩 ⁡ ( 0 , 0.1 ) \mathcal{N}(0,0.1) 13.43

### 4.3 xLSTM as Large Language Model

We culminate this study in large-scale language modeling experiments, testing the potential of xLSTM as an LLM. We therefore increase the amount of training data and train on 300B tokens from SlimPajama. The same number of tokens is used in, e.g., Mamba ( Gu & Dao, 2023 ) and Griffin ( De et al., 2024 ) . We compare xLSTM to RWKV-4, Llama, and Mamba – one method from each respective method class in Section 4.2 . We select RWKV-4 as RNN representative since for RWKV-5, RWKV-6 and HGRN2 a reasonable training precision setting has been found only after the training start of the 300B token experiments (see Appendix B.2 ). We train different model sizes (125M, 350M, 760M, 1.3B), test all models for length extrapolation capabilities and evaluate their performance on the validation set. We assess their performance on downstream tasks, test their performance in language modeling on 571 text domains of the PALOMA benchmark, and, finally, investigate their scaling law behavior.

##### Sequence Length Extrapolation.

Firstly, we test the sequence length extrapolation for 1.3B-sized, large models of xLSTM, RWKV-4, Llama, and Mamba. All models are trained on context length 2048, and then tested for context lengths up to 16384. See Figure 7 for the results. In contrast to other methods, xLSTM models maintain low perplexities for longer contexts.

Model SlimPajama (300B) ppl ↓ \downarrow at 16k Llama 337.83 Mamba 14.00 RWKV-4 13.75 xLSTM[7:1] 8.92 xLSTM[1:0] 9.01

##### Validation Perplexity and Downstream Tasks.

Secondly, for all model sizes, we evaluate the performance of xLSTM, RWKV-4, Llama, and Mamba models on the SlimPajama validation set for next token prediction and on downstream tasks that measure common sense reasoning. The third column of Table 3 lists the validation set perplexities of different methods. Both xLSTM[1:0] and xLSTM[7:1] are the best models for all model sizes with respect to the validation set perplexity. The other columns of Table 3 provide the performance on downstream tasks. In the vast majority of tasks and across all model sizes xLSTM is the best method — only on the ARC task Mamba is in some cases the best method. For details see Appendix B.3 .

Model #Params M SlimPajama (300B) ppl ↓ \downarrow LAMBADA ppl ↓ \downarrow LAMBADA acc ↑ \uparrow HellaSwag acc ↑ \uparrow PIQA acc ↑ \uparrow ARC-E acc ↑ \uparrow ARC-C acc ↑ \uparrow WinoGrande acc ↑ \uparrow Average acc ↑ \uparrow 125M RWKV-4 169.4 16.66 54.72 23.77 34.03 66.00 47.94 24.06 50.91 41.12 Llama 162.2 15.89 39.21 31.54 34.09 65.45 45.33 23.63 50.67 41.78 Mamba 167.8 15.08 27.76 34.14 36.47 66.76 48.86 24.40 51.14 43.63 xLSTM[1:0] 163.8 14.63 25.98 36.52 36.74 65.61 47.81 24.83 51.85 43.89 xLSTM[7:1] 163.7 14.60 26.59 36.08 36.75 66.87 48.32 25.26 51.70 44.16 350M RWKV-4 430.5 12.62 21.57 36.62 42.47 69.42 54.46 25.43 51.22 46.60 Llama 406.6 12.19 15.73 44.19 44.45 69.15 52.23 26.28 53.59 48.32 Mamba 423.1 11.64 12.83 46.24 47.55 69.70 55.47 27.56 54.30 50.14 xLSTM[1:0] 409.3 11.31 11.49 49.33 48.06 69.59 55.72 26.62 54.38 50.62 xLSTM[7:1] 408.4 11.37 12.11 47.74 47.89 71.16 56.61 27.82 53.28 50.75 760M RWKV-4 891.0 10.55 10.98 47.43 52.29 72.69 58.84 28.84 55.41 52.58 Llama 834.1 10.60 9.90 51.41 52.16 70.95 56.48 28.75 56.67 52.74 Mamba 870.5 10.24 9.24 50.84 53.97 71.16 60.44 29.78 56.99 53.86 xLSTM[1:0] 840.4 9.86 8.09 54.78 55.72 72.69 62.75 32.59 58.17 56.12 xLSTM[7:1] 839.7 9.91 8.07 55.27 56.12 72.74 61.36 29.61 56.43 55.26 1.3B RWKV-4 1515.2 9.83 9.84 49.78 56.20 74.70 61.83 30.63 55.56 54.78 Llama 1420.4 9.44 7.23 57.44 57.81 73.12 62.79 31.74 59.04 56.99 Mamba 1475.3 9.14 7.41 55.64 60.45 74.43 66.12 33.70 60.14 58.41 xLSTM[1:0] 1422.6 8.89 6.86 57.83 60.91 74.59 64.31 32.59 60.62 58.48 xLSTM[7:1] 1420.1 9.00 7.04 56.69 60.26 74.92 65.11 32.34 59.27 58.10

Model #Params M C4 MC4 EN Wikitext 103 Penn Treebank Red Pajama Refined Web Dolma M2D2 S2ORC M2D2 Wikipedia C4 Domains Dolma Subreddits Dolma Coding Average 125M RWKV-4 169.4 26.25 22.33 29.18 38.45 8.99 32.47 17.04 23.86 21.42 22.68 37.08 5.12 23.74 Llama 162.2 24.64 17.23 23.16 31.56 8.26 29.15 15.10 19.71 20.41 21.45 36.73 3.61 20.92 Mamba 167.8 23.12 17.04 22.49 30.63 7.96 27.73 14.60 19.38 19.36 20.14 34.32 3.77 20.05 xLSTM[1:0] 163.8 22.54 16.32 21.98 30.47 7.80 27.21 14.35 19.02 19.04 19.65 34.15 3.64 19.68 xLSTM[7:1] 163.7 22.39 16.13 21.47 30.01 7.75 26.91 14.13 18.6 18.84 19.52 33.9 3.59 19.44 350M RWKV-4 430.5 19.55 15.82 19.64 27.58 6.97 24.28 12.94 17.59 15.96 16.98 29.40 3.90 17.55 Llama 406.6 18.38 13.28 16.41 21.82 6.56 22.09 11.76 15.05 15.25 15.99 28.30 3.12 15.67 Mamba 423.1 17.33 13.05 16.11 22.24 6.34 21.04 11.42 14.83 14.53 15.16 27.02 3.20 15.19 xLSTM[1:0] 409.3 17.01 12.55 15.17 22.51 6.20 20.66 11.16 14.44 14.27 14.85 26.70 3.08 14.88 xLSTM[7:1] 408.4 16.98 12.68 15.43 21.86 6.23 20.70 11.22 14.62 14.30 14.85 26.61 3.11 14.88 760M RWKV-4 891.0 15.51 12.76 14.84 21.39 5.91 19.28 10.70 14.27 13.04 13.68 24.22 3.32 14.08 Llama 834.1 15.75 11.59 13.47 18.33 5.82 19.04 10.33 13.00 13.05 13.76 24.80 2.90 13.49 Mamba 870.5 15.08 11.54 13.47 19.34 5.69 18.43 10.15 13.05 12.62 13.25 23.94 2.99 13.30 xLSTM[1:0] 840.4 14.60 11.03 12.61 17.74 5.52 17.87 9.85 12.50 12.20 12.81 23.46 2.87 12.76 xLSTM[7:1] 839.7 14.72 11.11 12.68 17.61 5.55 18.01 9.87 12.59 12.25 12.89 23.43 2.88 12.80 1.3B RWKV-4 1515.2 14.51 12.04 13.73 19.37 5.62 18.25 10.11 13.46 12.10 12.87 22.85 3.25 13.18 Llama 1420.4 13.93 10.44 11.74 15.92 5.29 17.03 9.35 11.61 11.53 12.24 22.63 2.74 12.04 Mamba 1475.3 13.35 10.40 11.76 16.65 5.21 16.50 9.17 11.73 11.18 11.83 21.43 2.83 11.84 xLSTM[1:0] 1422.6 13.13 10.09 11.41 15.92 5.10 16.25 9.01 11.43 10.95 11.60 21.29 2.73 11.58 xLSTM[7:1] 1420.1 13.31 10.21 11.32 16.00 5.16 16.48 9.11 11.61 11.10 11.76 21.50 2.75 11.69

##### Performance on PALOMA Language Tasks.

Thirdly, for all model sizes, we test the next token prediction performance of xLSTM, RWKV-4, Llama, and Mamba models on PALOMA language tasks ( Magnusson et al., 2023 ) . We measure the performance by the perplexity for next token prediction on 571 text domains, which range from nytimes.com to r/depression on Reddit. Table 4 shows token prediction perplexity grouped into language modeling (first seven columns) and fine-grained domain benchmarks (last 5 columns). xLSTM[1:0] performs better than xLSTM[7:1] on these language tasks. xLSTM[1:0] has in 568 out of 571 (99.5%) text domains a lower perplexity than Mamba, in 486 out of 571 (85.1%) a lower perplexity than Llama, in 570 out of 571 (99.8%) a lower perplexity than RWKV-4. For details see Appendix B.3 .

##### Scaling Laws.

Fourthly, we assess the power-law scaling behavior, which allows to extrapolate the performance to larger model sizes ( Kaplan et al., 2020 ; Brown et al., 2020 ) . Figure 8 presents the scaling behavior. All models share a similar scaling behavior but with different offsets. RWKV-4 performs worst, followed by Llama and Mamba. xLSTM is better than Mamba with a similar margin to Mamba as Mamba has to Llama. The scaling behavior indicates that for larger models xLSTM will continue to perform favourable compared to Transformers and State-Space models.

Generation Times and Maximal Throughput. Finally, we measure the text generation time in Figure 9 and the maximal throughput in Figure 9 (left) for our xLSTM variants at 1.3B scale. We compare against similar sized Mamba, Llama and RWKV implementions from HuggingFace, including a static key-value cache for the Llama model. At the time of the experiments, both full cache compilation of the Transformer Model and compilation of the Mamba model with torch.compile did not work. For the text generation experiments all of the models are tested at batch size 1 and pre-fill 16. This pre-fill should be maximally favorable for the Transformer. Figure 9 shows the linear scaling of the xLSTM and the other recurrent models Mamba and RWKV-4 compared to the quadratic scaling of Llama. For the decoding throughput we measure different batch sizes and prefill for the Llama model. Figure 9 (right) shows that xLSTM can use much higher batch sizes than Llama due to its constant memory and thus achieves the highest throughput.

## 5 Limitations

(i) In contrast to mLSTM, memory mixing of the sLSTM prohibits parallelizable operations, and therefore does not allow a fast parallel implementation. Nevertheless, we developed a fast CUDA kernel for sLSTM, which is currently less than two times slower than our parallel mLSTM implementation. (ii) The CUDA kernels for mLSTM are not optimized, and therefore the current implementation is about 4 times slower than FlashAttention or the scan used in Mamba. Faster CUDA kernels could be obtained in the vein of FlashAttention. (iii) The matrix memory of mLSTM has high computation complexity since d × d d\times d matrices must be processed. Still, the memory update and retrieval does not use parameters and can be parallelized using standard matrix operations, therefore the wall clock time overhead due to the complex memory is minor. (iv) The initialization of the forget gates must be chosen carefully. (v) Since the matrix memory is independent of the sequence length, increasing the sequence length might overload the memory for longer context sizes. Still, this does not appear to be a limitation for contexts up to 16k, see Section 4.3 . (vi) Due to the expensive computational load for large language experiments, we did neither fully optimize the architecture nor the hyperparameters, especially for larger xLSTM architectures. We anticipate that an extensive optimization process is needed for xLSTM to reach its full potential.

## 6 Conclusion