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

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

### A.4 Detailed Block Structure

## Appendix B Experiments

##### Training Setup.

For all experiments, we use Python 1 1 1 https://python.org 3.11 with PyTorch 2.2.0 2 2 2 https://pytorch.org , and CUDA 12.1 3 3 3 https://docs.nvidia.com/cuda/archive/12.1.0/ on NVIDIA A100 GPUs. We developed and trained all our models and baselines over the course of three months on a cluster with 128 nodes of eight NVIDIA A100 GPUs each. More than 95% of this compute were used for the Language Modeling experiments in Sections 4.2 and 4.3 .

##### Nearest Neighbor Search Task.

For this auxiliary task, we use randomly sampled feature vectors of dimension 2 and unit norm. The attached value is a uniformly distributed random number from [ 0 , 1 ] [0,1] , leading to inputs vectors of dimension 3. The first feature vector serves as search key, with the first value being ignored. Then the model has to predict the value of the nearest neighbor so far in the sequence. We train on 8192 sequences of context length up to 64 (uniformly sampled) and validate on 8192 different samples. All models have two blocks and embedding dimension 128. We use a dropout of 0.1, 10% linear warm-up steps and cosine decay to 1e-7 for 100k total training steps. We sweep over learning rates 1e-4, 1e-3, 1e-2, 1e-1 and 5 seeds each. The reported values in Figure 2 are mean values for the best learning rate and 99% confidence intervals. Note that LSTM requires very high learning rates, whereas Transformers (Llama) perform best at the smallest learning rate. The xLSTM[0:1] reaches similar performance across all learning rates.

##### Wikitext-103 Rare Token Prediction.

For this exemplary experiment on rare token prediction, we trained 125M-sized models on Wikitext-103 ( Merity et al., 2017 ) . All models have an embedding dimension of 768 in a post up-projection structure of 12 residual blocks. The Transformer model (Llama) uses Multi-Head Attention, for what is called LSTM the Multi-Head Attention is replaced by an LSTM and the xLSTM[1:0] contains mLSTM layers with matrix memory. Models were trained with maximum learning rate 1e-3, 4k steps linear warm-up and cosine decay for in total 50k steps, using a batch size of 256 and context length of 512. We use the validation perplexity as a stopping criterion and evaluate on the test set.

### B.1 Synthetic Tasks and Long Range Arena

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

We evaluate xLSTM on a suite of formal language tasks to test its exponential gating and memory mixing mechanism.

Formal languages provide a framework to probe the generalization capabilities of models. They allow to specifically test different expressivity levels, e.g. along the Chomsky hierarchy. Typical language model architectures do not necessarily fit perfectly in these hierarchies ( Delétang et al., 2023 ) — nevertheless these languages allow to illustrate differences in generalization expressivity between different architectures. Our evaluation tasks are heavily based on the work of Delétang et al. (2023) .

##### Experiment Setup.

The different formal language tasks in the experiment (see individual tasks description below) encompass different levels of the Chomsky hierarchy as well as additional counting and memory-focused tasks. We use different lengths per sample, which allows us to validate in a length extrapolation setting. We train on a varying task length up to 40. The evaluation is done for task lengths between 40 and 256 as we are only interested in the “task generalization capabilities“ of the models.

In all experiments, we use two blocks (or layers for the pure LSTM) for all models. We compare Llama, Mamba, Retention, Hyena, RWKV-4, RWKV-5, RWKV-6, LSTM, xLSTM[0:1], xLSTM[1:0] and xLSTM[1:1]. The sLSTM block is used without a convolution and with normal weight initialization. LSTM (Block) refers to an architecture where a vanilla LSTM is used instead of self-attention inside a Transformer block.

All models are trained with 3 different learning rates (1e-2, 1e-3, 1e-4), each with two seeds. Batch size is 256 — cosine annealing (min lr: 1e-5) with 10% warm-up steps is applied. We use AdamW ( Loshchilov & Hutter, 2019 ) ( β 1 = 0.9 \beta_{1}=0.9 , β 2 = 0.99 \beta_{2}=0.99 ) and a weight decay of 0.1 for training. In each experiment we train for 100k steps — the samples are generated randomly, however, all experiments are trained and evaluated on the same samples.

##### Additional Formal Language Results.

Figure 12 showcases supplementary results on formal language task, detailing tasks where no model attained a minimum scaled accuracy of 0.3. Although no model achieves proper extrapolation of the task to a larger context length, xLSTM performs best among the evaluated models.

##### Individual Task Description.

The majority of tasks are based on Delétang et al. (2023) . We provide the vocabulary size | V | |V| and the random accuracy s r ​ a ​ n ​ d s_{rand} (for accuracy scaling), used in the evaluation. As we evaluate different task lengths each task has a padding token which is used to pad the sequence to the given context length. In Listing 1 there is an example for each task.

• Bucket Sort Given a string of tokens of a sorted alphabet, compute the sorted string. | V | = 11 s rand = 1 | V | − 1 |V|=11\quad s_{\text{rand}}=\frac{1}{|V|-1}

• Cycle Nav Given a string of “movement tokens” ( + 1 +1 , − 1 -1 , STAY) compute the end position of the agent with start position 0. The position must be computed modulo the maximum position. | V | = 9 s rand = 1 | V | − 4 |V|=9\quad s_{\text{rand}}=\frac{1}{|V|-4}

• Even Pairs Given a binary string of a and b tokens, compute whether the number of ab and ba is even. This task can be solved by checking if the first and last token of the string are equal. | V | = 3 s rand = 0.5 |V|=3\quad s_{\text{rand}}=0.5

• Majority Given a string of tokens, compute the token that occurred most often in the sequence. | V | = 64 s rand = 1 | V | − 1 |V|=64\quad s_{\text{rand}}=\frac{1}{|V|-1}

• Majority Count Given a string of tokens of an ordered alphabet. Compute the count of the token that occurred most often in the sequence. If the count exceeds the vocab size, the highest vocab token should be outputted. | V | = 64 s rand = 1 | V | − 1 |V|=64\quad s_{\text{rand}}=\frac{1}{|V|-1}

• Missing Duplicate Given a string of tokens. The string is repeated but one of the tokens is masked in the repetition. Output the token that is masked. | V | = 11 s rand = 1 | V | − 2 |V|=11\quad s_{\text{rand}}=\frac{1}{|V|-2}

• Mod Arithmetic (w/o Brackets) Calculate the result — modulo the max number — of the arithmetic operations in the context. The maximum number is the vocabulary size minus the number of special tokens (+,-,*,=, [PAD]). | V | = 10 s rand = 1 | V | − 5 |V|=10\quad s_{\text{rand}}=\frac{1}{|V|-5}

• Mod Arithmetic (w Brackets) Calculate the result — modulo the maximum number — of the arithmetic operations in the context. The maximum number is vocabulary size minus the number of special tokens (+,-,*,=,(,), [PAD]). | V | = 12 s rand = 1 | V | − 7 |V|=12\quad s_{\text{rand}}=\frac{1}{|V|-7}

• Odds First An string of tokens t 1 , t 2 , t 3 , … ​ t n t_{1},t_{2},t_{3},...t_{n} is given. Output all tokens with and odd index ( t 1 , t 3 , … t_{1},t_{3},... ) then the token with an even index ( t 2 t_{2} , t 4 t_{4} ,..) . Apart from that keep the ordering of the initial string. | V | = 12 s rand = 1 | V | − 2 |V|=12\quad s_{\text{rand}}=\frac{1}{|V|-2}

• Parity Given a binary string of a and b tokens, compute if the number of b ‘s is even. If the number is even output a otherwise b . This is equivalent to sequentially calculating the half-adder sum. | V | = 3 s rand = 0.5 |V|=3\quad s_{\text{rand}}=0.5

• Repetition Given a string of tokens — repeat it. | V | = 12 s rand = 1 | V | − 2 |V|=12\quad s_{\text{rand}}=\frac{1}{|V|-2}

• Reverse String Given a string of tokens — repeat it in reverse order. | V | = 12 s rand = 1 | V | − 2 |V|=12\quad s_{\text{rand}}=\frac{1}{|V|-2}

• Stack Manipulation An initial stack content is given, followed by a sequence of push and pop operations. Compute the stack content after the operations | V | = 11 s rand = 1 ⌊ | V | − 3 2 ⌋ |V|=11\quad s_{\text{rand}}=\frac{1}{\lfloor\frac{|V|-3}{2}\rfloor}

• Set Given a string of tokens, compute the ordered set of the tokens. Keep the ordering so that tokens that occurred first are also outputted first. | V | = 128 s rand = 1 | V | − 2 |V|=128\quad s_{\text{rand}}=\frac{1}{|V|-2}

• Solve Equation Given is an equation with the operators {+,-,*,=,(,)}, number, and an unknown variable x. Compute the value of the variable modulo the max number. The maximum number is vocabulary size minus the number of special tokens (+,-,*,=,(,), [PAD], [ACT]). | V | = 14 s rand = 1 | V | − 9 |V|=14\quad s_{\text{rand}}=\frac{1}{|V|-9}

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

We test the memory capacity of xLSTM with the Multi-Query Associative Recall task proposed by Arora et al. (2023) . Figure 5 illustrates the basic task setup.

##### Why Multi-Query Associative Recall for Memory Tests of LLM Architectures.

Associative Recall (AR), the ability to retrieve a specific value (information) associated with a given key (information), constitutes a key capability for LLM to perform well ( Poli et al., 2024 ; Arora et al., 2023 ; Olsson et al., 2022 ) . Especially its quality of in-context learning seems to be strongly connected to this capability ( Olsson et al., 2022 ) . Arora et al. (2023) attribute performance gaps between early non-Transformer and Transformer language models specifically to performance gaps in associative recall. They argue that prior AR evaluations fall short of capturing these differences and propose MQAR, which can show the AR performance differences that translate to performance differences in language modeling performance. Hence, MQAR is especially suitable to analyze the memory capacity of LLM. Transformer (e.g. Llama) models can be seen as the gold standard for this task as their memory is exponential in the coding dimension ( Ramsauer et al., 2021 ) .

##### Experiment Setup.

There are two relevant variables that determine different experimental setups. (1) Context Length (CL) : Length of the sequence of one sample — this influences the distances between the key–value definition and the recall. (2) Number Key-Value Pairs (KV) : Influences how many key-value pairs the model needs to keep track of. The vocabulary size is always 8192.

In all experiments, we use two blocks (or layers for the pure LSTM) for all models. LSTM (Block) model refers to an architecture where a vanilla LSTM is used instead of self-attention inside a Transformer block.

For each task setup, we train each model with 4 different learning rates (batch size > 24: {1e-2, 2.15e-3, 4.6e-4, 1e-4}, batch size 24: {1e-3, 2.2e-4, 5e-5, 1e-5}). The batch size (BS) changes depending on the context length (CL) (CL=64/128: BS=512; CL=256: BS=256; CL=756: BS=128; CL=1024: BS=96; CL=2048: BS=24). We vary the embedding dimension (Model Dim) between different experiments – different numbers of heads are used accordingly. For each experiment, we generate 100,000 training samples (validation: 3,000 samples) and train for 64 epochs. We apply cosine annealing (min lr: 1e-4 and 1e-5) with 10% warm-up steps. We use AdamW ( Loshchilov & Hutter, 2019 ) and a weight decay of 0.1 for training.

We conduct three different experiments:

• MQAR-Experiment 1 evaluates, in the same fashion as Arora et al. (2023) , a variety of models (Llama, Mamba, Mamba (noWT) - i.e. without weight tying, Retention, Hyena, H3, RWKV-4, RWKV-5, RWKV-6, LSTM, LSTM (Block), xLSTM[0:1], xLSTM[1:0] and xLSTM[1:1]) on increasing task difficulty by increasing the context length and number of key-value pairs simultaneously. We benchmark three parameter settings: CL,KV={(64,4),(128,8),(256,16)}.

• MQAR-Experiment 2 increases the task difficulty notably and goes beyond previous evaluations on this task. We individually scale the context length (CL={756, 1024, 2048}) and the key-value pairs (KV={48, 96, 256}) and evaluate all combinations. This experiment especially probes the memory capacity because the number of key-value pairs is high. To reduce the computational burden we only evaluate models that perform flawlessly in Experiment 1 — additionally we evaluate Transformer only in the hardest setting (CL=2048) as sanity check, because no performance decrease is expected.

• MQAR-Experiment 3 analyzes whether the AR capability learned on a certain context length extrapolates to bigger context lengths. For each KV setting of Experiment 2, we use the models (we select the 3 biggest model dimensions) trained on CL=2048 and evaluate bigger context lengths (CL={4096, 6144, 8192}).

##### Extended Results.

The result of Experiment 1 can be found in Figure 14 . In accordance to the results of Arora et al. (2023) H3, Hyena, RWKV-4 fail to solve the task with a smaller model dimension. In contrast, xLSTM[1:1], xLSTM[1:0], Mamba, RWKV-5 and RWKV-6 are able to solve these settings for all model dimensions. The comparison of xLSTM[0:1] with both original LSTM variants indicates that the exponential gating mechanism improves the AR capabilities of the model. However, both fall short because of the reduced memory capacity compared to xLSTM[1:1] and xLSTM[1:0].

The results of Experiment 2 are presented in Figure 15 . Scaling the context length has a low impact on the performance of the models. However, while xLSTM[1:1] and xLSTM[1:0] show no clear decay, both RWKV variants slightly, but consistently lose performance with increasing context lengths. The varying number of key-value pairs, which mainly probes the memory capacity of the non-Transformer models, has a more notable impact across all models. RWKV-5 seems to outperform RWKV-6. The latter fails to learn the task at all in some KV=256 settings. Overall xLSTM[1:1] is the best-performing non-Transformer model — suggesting that it provides enhanced memory capacity, also in long contexts.

Figure 16 shows the extrapolation results from Experiment 3. For xLSTM[1:1], xLSTM[1:0], and Mamba the model performance does not change in the extrapolation setting. The RWKV models (especially RWKV5) degrade slightly with increasing context length. xLSTM[1:1] performs best, as it maintains its superior performance of Experiment 2.

#### B.1.3 Test of xLSTM’s Long Range Capabilities on the Long Range Arena.

We assess the performance of xLSTM across tasks in the Long Range Arena benchmark ( Tay et al., 2021 ) , examining its ability to effectively handle longer context lengths and diverse data types.

Our experiments on Long Range Arena benchmark are composed of five tasks: • Retrieval : The task is to predict if two documents have a citation link. The dataset of text documents is derived from the ACL Anthology Network ( Radev et al., 2009 ) .

• ListOps : This is a set of modular arithmetic tasks including brackets and lists of numbers, using the operations MIN , MAX , MEDIAN and SUMMOD (modular sum). A particular example is: [MAX 4 3 [MIN 2 3 ] 1 0 [MEDIAN 1 5 8 9, 2]] → \xrightarrow{} 5

• Image : This task is based on a version of the CIFAR dataset ( Krizhevsky, 2009 ) , where images are transformed to a sequence of pixels and this sequence has to be classified into the usual CIFAR classes. We test both a gray-scale (G-Image) and RGB (RGB-Image) version of this dataset, as Orvieto et al. (2023) uses colored images contrary to the standard setup.

• Pathfinder : The input for this task is a 32x32 gray-scale image, given as pixel sequence, with two dots and several curved lines on it. The task is to predict if the two dots are connected by any of the lines ( Linsley et al., 2018 ) .