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

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

We omit the Text classification task ( Maas et al., 2011 ) , as the language modeling experiments already test this kind of data, and the Pathfinder-X version of Pathfinder .

##### Experiment Setup.

The architectures that are tested in this experiment comprise Llama, Mamba, LSTM, RWKV-4, and xLSTM. LSTM (Block) refers to an architecture where a vanilla LSTM is used inside a post up-projection block (like Transformer with attention replaced by LSTM). For xLSTM we choose the best performing of xLSTM[0:1] or xLSTM[1:0] on the validation set, specifically the former for the Image tasks and the latter for all other ones.

We use the hyperparameter settings of the S5 model ( Smith et al., 2022 ) and Linear Recurrent Unit model ( Orvieto et al., 2023 ) , with additional hyperparamter search on learning rates and schedulers for all models. We use two different schedulers: Linear Warm-up Cosine Annealing and Linear Warm-up Cosine Annealing with Restarts. Both learning rate schedulers were evaluated with learning rates of 1e-3, 6e-4 and 1e-4. For the second scheduler, the number of restarts ( R R ) is set to 3. The model hyperparameters for each dataset are displayed in Table 5 .

##### Results.

Table 6 shows the result of experiments on the Long Range Arena benchmark. xLSTM demonstrates consistent strong performance on all of the tasks, suggesting that the proposed architecture is remarkably efficient in handling different aspects of long context problems.

Retrieval acc ↑ \uparrow ListOps acc ↑ \uparrow Pathfinder acc ↑ \uparrow G-Image acc ↑ \uparrow RGB-Image acc ↑ \uparrow Ranking acc ↑ \uparrow Random Baseline 0.500 0.100 0.500 0.100 0.100 Llama 0.845 0.379 0.887 0.541 0.629 5.2 Mamba 0.902 0.325 0.992 0.689 0.765 2.2 RWKV-4 0.898 0.389 0.914 0.691 0.757 3.0 LSTM X 0.275 X 0.675 0.718 5.4 LSTM (Block) 0.880 0.495 X 0.690 0.756 3.4 xLSTM 0.906 0.411 0.919 0.695 0.761 1.6

### B.2 Method Comparison and Ablation Study on SlimPajama (15B)

##### General Training Procedure.

We tokenize our datasets using the HuggingFace GPT-2 tokenizer ( Radford et al., 2019 ; Brown et al., 2020 ) 6 6 6 https://huggingface.co/docs/transformers/en/model_doc/gpt2 and use this tokenizer for all models in this paper. In general, we try to follow Brown et al. (2020) for the general training setup, i.e. we choose context length 2048 and batch sizes 256 or 512 for our models. We use the AdamW ( Loshchilov & Hutter, 2019 ) optimizer with beta parameters ( β 1 \beta_{1} , β 2 \beta_{2} )=(0.9, 0.95) and an epsilon parameter of 1e-5. As learning rate scheduler we use a linear warm-up with 750 steps and cosine decay to 10% of the peak learning rate. We apply a weight decay of 0.1 to all our models and always exclude the token embedding matrix from weight decay. If not specified otherwise, we do not tie the weights of the token embedding and the language model head. For parallelization, we use PyTorch FSDP in SHARD_GRAD_OP mode with mixed precision in bfloat16 , where applicable. For small models we use NO_SHARD . We keep the weights in float32 and reduce the gradients across GPUs in float32 . We use torch.compile to speed up models, except for Transformer models as their training curves did not match the non-compiled versions. For xLSTM[7:1], we use positions [3, 5, 7, 40, 42, 44] for sLSTM-based blocks, except for the 125M size, where we use [3, 20] (this is actually a [11:1] ratio). We do not use any positional encoding for our xLSTM models.

Model EmbeddingDim #Blocks #Heads/HeadDim #Params M Peak LR (15B) Peak LR (300B) 125M RWKV-5 768 12 - 176.5 3e-3 - RWKV-6 768 12 - 173.6 3e-3 - HGRN2 768 12 - 162.2 3e-3 - RWKV-4 768 12 - 169.4 3e-3 6e-4 Llama 768 12 12 / 64 162.2 3e-3 3e-3 Mamba 768 24 - 167.8 3e-3 3e-3 xLSTM 768 24 4 / 384 163.8 3e-3 1.5e-3 350M RKWV-5 1024 24 - 455.7 1e-3 - RWKV-6 1024 24 - 441.6 1e-3 - HGRN2 1024 24 - 411.4 1e-3 - RWKV-4 1024 24 - 430.5 1e-3 4e-4 Llama 1024 24 16 / 64 406.6 1.5e-3 1.5e-3 Mamba 1024 48 - 423.1 1.5e-3 1.5e-3 xLSTM 1024 48 4 / 512 409.3 1e-3 7.5e-4 760M RWKV-5 1536 24 - 947.8 9e-4 - RWKV-6 1536 24 - 907.7 9e-4 - HGRN2 1536 24 - 834.2 9e-4 - RWKV-4 1536 24 - 891.0 2e-3 2.5e-4 Llama 1536 24 16 / 96 834.1 1.25e-3 1.25e-3 Mamba 1536 48 - 870.5 1.25e-3 1.25e-3 xLSTM 1536 48 4 / 768 840.4 9e-4 6.25e-4 1.3B RWKV-5 2048 24 - 1616.0 9e-4 - RWKV-6 2048 24 - 1537.5 9e-4 - HGRN2 2048 24 - 1439.4 9e-4 - RWKV-4 2048 24 - 1515.2 1e-3 2e-4 Llama 2048 24 16 / 128 1420.4 1e-3 1e-3 Mamba 2048 48 - 1475.3 1e-3 1e-3 xLSTM 2048 48 4 / 1024 1422.6 9e-4 5e-4 2.7B RWKV-5 2048 24 - 3194.7 8e-4 - RWKV-6 2048 24 - 3021.9 8e-4 - HGRN2 2048 24 - 2795.4 8e-4 - RWKV-4 2560 32 - 2984.8 8e-4 - Llama 2560 32 32 / 80 2779.5 8e-4 - Mamba 2560 64 - 2897.2 8e-4 - xLSTM 2560 64 4 / 1280 2788.3 8e-4 -

##### Details on Comparison to Other Methods.

For the model comparison on 15B training tokens of SlimPajama we train all models with context length 2048 and batch size 256. We use a peak learning rate of 1e-3 for all models for comparability. The learning rate decays over 30k training steps. The models are compared after one epoch at training step 28170. As model implementations we use the original repositories’ code for Mamba ( Gu & Dao, 2023 ) 7 7 7 https://github.com/state-spaces/mamba , RWKV-5, RWKV-6 ( Peng et al., 2024 ) 8 8 8 https://github.com/BlinkDL/RWKV-LM/ . For RWKV-4 we use a cleaned and validated re-implementation based on the original repo and kernels ( Peng et al., 2023 ) . In our RWKV-4 implementation we enable weight decay on all parameters except biases, the token embedding weight and all LayerNorm weights. For HGRN ( Qin et al., 2023 ) , GLA ( Yang et al., 2023 ) , HGRN2 ( Qin et al., 2024 ) we use the a re-implementation flash-linear-attention ( Yang & Zhang, 2024 ) by the authors of GLA ( Yang et al., 2023 ; Yang & Zhang, 2024 ) 9 9 9 https://github.com/sustcsonglin/flash-linear-attention . For GPT-3 and Llama-like Transformers, we use our own implementations based on PyTorch. Note that for all xLSTMs, Transformers, Mamba and RWKV-4, we use Mixed Precision training with bfloat16 and weights in float32 precision. Following the general training procedure we use torch.compile for all models, except for Transformers and models using the flash-linear-attention library because of compilation problems.

As RWKV-6 performs worse than RWKV-5, we also train a model with peak learning rate 4e-4, as reported in the original repository for 350M parameter models 10 10 10 https://github.com/BlinkDL/RWKV-LM/blob/64b7fe4c66fcf7da37019630268075b0558f6dc5/RWKV-v5/train.py#L44 . This model reaches a perplexity of 16.38, worse than the 15.03 for the standard peak learning rate 1e-3 as reported in Table 1 . Similarly, we tested the repository learning rates for other model sizes and all performed worse than the ones we also use for xLSTM (see Table 7 ).

##### Details on Training Precision for Baselines.

For models from flash-linear-attention and RWKV-5/6 models we found that PyTorch automatic mixed precision training did not work, but casting the model weights to float32 initially with FSDP parameter precision bfloat16 led to a working configuration. In this setting, models perform better than in full bfloat16 training, where the weights are casted to bfloat16 initially as well. Full float32 training did not work because of the custom kernels that require bfloat16.

##### General Details on Ablation Studies.

We follow our general training procedure and train all models with context length 2048, batch size 256 and peak learning rate 1e-3. We report perplexity values on the validation set.

##### Additional Ablation Study on Matrix Memory.

As default block configuration we use the mLSTM in the pre up-projection block (see Figure 11 ) and the sLSTM in the post up-projection block (see Figure 10 ). In this experiment we study combination of mLSTM with different block variants using the xLSTM[1:0] architecture. We compare the mLSTM in a post up-projection block (see Figure 3 and 10 ) with ReLU 2 \text{ReLU}^{2} activation function and non-gated feed-forward network to mLSTM in a pre up-projection block with and without a dimension-wise causal convolution. Table 8 shows that the matrix memory benefits from the pre up-projection block structure, and that the convolution within this block is important.

Model Details #Blocks Embedding Dim #Params M SlimPajama (15B) ppl ↓ \downarrow xLSTM[1:0] Post Up-Projection Block (ReLU2) 24 1024 430.4 13.90 Pre Up-Projection Block, No Convolution 48 1024 408.8 15.41 Pre Up-Projection Block, With Convolution 48 1024 409.3 13.43

##### Details on new xLSTM Components Ablation Study.

In Table 2 (top), we show our modifications to the vanilla LSTM that transform the vanilla LSTM into the xLSTM. We start with a large default PyTorch LSTM with 24 layers and 1536 hidden size. Due to a lack of skip-connections and LayerNorms, vanilla LSTMs of this size are not trainable. We then add skip-connections and pre-LayerNorms before each LSTM layer corresponding to a residual architecture. This enables training for LSTMs at this scale. Replacing every second LSTM layer by a non-gated feed-forward network with GeLU activation function (similar to Vaswani et al. ), which corresponds to the post up-projection backbone (see Figure 3 ) further boosts performance. Adding Exponential Gating to this architecture yields the sLSTM as depicted in Figure 10 , with another large performance improvement. Finally, adding the best Matrix Memory variant found in Table 8 by replacing some sLSTM blocks with the mLSTM (see Figure 11 ) gives xLSTM[7:1] with the best performance.

##### Details on Gating Technique Ablation Study.

In Table 2 (bottom), we investigate the effect of trainable and input-dependent gates for mLSTM. The results show that, in contrast to other methods ( Katharopoulos et al., 2020 ; Sun et al., 2023 ; Qin et al., 2023 ; Katsch, 2023 ; Yang et al., 2023 ; Qin et al., 2024 ; Peng et al., 2024 ) , having the gates both learnable and input dependent gives the best results.

##### Details on Scaling Experiments.

We follow our general training procedure (see paragraph above) and train all models, including the 1.3B and 2.7B model sizes, with context length 2048 and batch size 256. We use the peak learning rates from Table 7 . For Llama and Mamba we use the learning rates reported by Gu & Dao (2023) .

### B.3 xLSTM Large Language Models – SlimPajama300B

##### General Training Procedure.

We use the same general training procedure as in Section B.2 with peak learning rates from Table 7 . For Llama and Mamba we use the learning rates reported by Gu & Dao (2023) . All models are trained with context length 2048. The 125M, 350M and 760M models are trained with batch size 256 for 600k training steps, whereas the 1.3B models are trained with batch size 512 for 300k training steps. We keep the same learning rate scheduler across all models.

##### Details on Downstream Evaluation.

We use the LM Evaluation Harness from EleutherAI ( Sutawika et al., 2023 ) for evaluating the following tasks that measure common sense reasoning: LAMBADA (OpenAI version in LM Evaluation Harness) ( Paperno et al., 2016 ) , HellaSwag ( Zellers et al., 2019 ) , PIQA ( Bisk et al., 2020 ) , ARC-challenge, ARC-easy ( Clark et al., 2018 ) , WinoGrande ( Sakaguchi et al., 2021 ) . This selection of downstream tasks is also used in previous work by Gu & Dao (2023) .

Following Gu & Dao (2023) , we report accuracy for LAMADA, WinoGrande, PIQA, and ARC-easy, and accuracy normalized by sequence length for HellaSwag and ARC-challenge.

We evaluate all models in full float32 , full bfloat16 and bfloat16 Mixed Precision with weights in float32 . For each model we select the best value respectively.

##### Details on PALOMA.

We use 16 out of the 18 data sources of the PALOMA dataset ( Magnusson et al., 2023 ) . We use C4 ( Raffel et al., 2019 ) , MC4-EN ( Xue et al., 2021 ) , Wikitext-103 ( Merity et al., 2017 ) , PennTreebank ( Vadas & Curran, 2011 ) , RedPajama ( TogetherComputer, 2023 ) , Falcon Refinedweb (Refined Web) ( Penedo et al., 2023 ) , Dolma v1.5 ( Soldaini et al., 2023 ) , M2D2 S2ORC, M2D2 Wikipedia ( Reid et al., 2022 ) , C4-100-Domains (C4 Domains) ( Chronopoulou et al., 2022 ) , Dolma-100-Subreddits (Dolma Subreddits) ( Soldaini et al., 2023 ) , Dolma-100-Programming Languages (Dolma Coding) ( Soldaini et al., 2023 ; Kocetkov et al., 2022 ) , TwitterAAE ( Blodgett et al., 2016 ; Liang et al., 2023 ) , Manosphere Corpus ( Ribeiro et al., 2021 ) , GAB Corpus ( Zannettou et al., 2018 ) , 4CHAN Corpus ( Papasavva et al., 2020 ) . We leave out ThePile ( Gao et al., 2021 ) and ICE ( Greenbaum & Nelson, 1996 ) as they are not part of Paloma’s Huggingface dataset repository 11 11 11 https://huggingface.co/datasets/allenai/paloma . A detailed description of these datasets can be found in Magnusson et al. (2023, Table 2) . All models are evaluated in bfloat16 Mixed Precision.

Results on the data sources TwitterAAE, Manosphere, GAB and 4CHAN are reported in Table 9 and for each individual dataset the results are given in Section C .

Model #Params M Twitter AAE Manosphere 4CHAN GAB 125M RWKV-4 169.4 265.80 39.31 18.48 53.89 Llama 162.2 277.93 32.98 14.03 56.45 Mamba 167.8 258.17 32.14 14.01 51.58 xLSTM[1:0] 163.8 244.53 31.45 13.27 51.00 xLSTM[7:1] 163.7 248.51 30.90 13.45 50.25 350M RWKV-4 430.5 216.17 30.25 13.82 42.25 Llama 406.6 231.09 25.90 11.49 43.04 Mamba 423.1 202.88 25.24 11.60 40.78 xLSTM[1:0] 409.3 200.61 24.58 11.20 39.83 xLSTM[7:1] 408.4 206.25 24.73 11.31 39.86 760M RWKV-4 891.0 195.27 24.66 12.00 35.73 Llama 834.1 205.50 22.69 10.40 37.68 Mamba 793.2 182.74 22.58 10.47 36.25 xLSTM[1:0] 840.4 179.74 21.66 10.11 35.33 xLSTM[7:1] 839.7 180.19 21.78 10.22 34.89 1.3B RWKV-4 1515.2 174.87 23.51 11.34 33.18 Llama 1420.4 192.52 20.67 9.67 34.84 Mamba 1475.3 171.38 20.37 9.80 32.01 xLSTM[1:0] 1422.6 166.16 19.94 9.64 31.90 xLSTM[7:1] 1420.1 171.36 20.28 9.64 32.17

In order to evaluate the perplexity values on each data source, we split the text documents into sequences of length 2048, which corresponds to the pre-training context length of all models. For documents longer than 2048 tokens we split each document into non-overlapping input sequences. In this case for the last input sequence, we follow the LM Evaluation Harness and fill up the full 2048 token context window with previous tokens, but compute the perplexity only on the remaining tokens.

We compute the token perplexities per data source in Table 4 as the exponential of the negative loglikelihoods per domain weighted by the number of tokens per domain in that data source as it is defined in Magnusson et al. (2023, Equation 1)

## Appendix C Detailed Results on PALOMA Language Model Evaluation

We report the perplexity values on each of the 571 subdomains of PALOMA in Table 10 . Note that the aggregated perplexity values in Table 4 are not macro averages of the values shown in Table 10 .

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .