Title: xLSTM: Extended Long Short-Term Memory
Venue: NeurIPS 2024
Accepted: yes
Keywords: LSTM

LLM

Language modeling

NLP

Memory
Abstract: In the 1990s, the constant error carousel and gating were introduced as the central ideas of the Long Short-Term Memory (LSTM). Since then, LSTMs have stood the test of time and contributed to numerous deep learning success stories, in particular they constituted the first Large Language Models (LLMs). However, the advent of the Transformer technology with parallelizable self-attention at its core marked the dawn of a new era, outpacing LSTMs at scale.  We now raise a simple question: How far do we get in language modeling when scaling LSTMs to billions of parameters, leveraging the latest techniques from modern LLMs, but mitigating known limitations of LSTMs? Firstly, we introduce exponential gating with appropriate normalization and stabilization techniques. Secondly, we modify the LSTM memory structure, obtaining: (i) sLSTM with a scalar memory, a scalar update, and new memory mixing, (ii) mLSTM that is fully parallelizable with a matrix memory and a covariance update rule. Integrating these LSTM extensions into residual block backbones yields xLSTM blocks that are then residually stacked into xLSTM architectures. Exponential gating and modified memory structures boost xLSTM capabilities to perform favorably when compared to state-of-the-art Transformers and State Space Models, both in performance and scaling.

xLSTM: Extended Long Short-Term Memory

Abstract

In the 1990s, the constant error carousel and gating were introduced as the central ideas of the Long Short-Term Memory (LSTM). Since then, LSTMs have stood the test of time and contributed to numerous deep learning success stories, in particular they constituted the first Large Language Models (LLMs). However, the advent of the Transformer technology with parallelizable self-attention at its core marked the dawn of a new era, outpacing LSTMs at scale. We now raise a simple question: How far do we get in language modeling when scaling LSTMs to billions of parameters, leveraging the latest techniques from modern LLMs, but mitigating known limitations of LSTMs? Firstly, we introduce exponential gating with appropriate normalization and stabilization techniques. Secondly, we modify the LSTM memory structure, obtaining: (i) sLSTM with a scalar memory, a scalar update, and new memory mixing, (ii) mLSTM that is fully parallelizable with a matrix memory and a covariance update rule. Integrating these LSTM extensions into residual block backbones yields xLSTM blocks that are then residually stacked into xLSTM architectures. Exponential gating and modified memory structures boost xLSTM capabilities to perform favorably when compared to state-of-the-art Transformers and State Space Models, both in performance and scaling.

1 Introduction

The Long Short-Term Memory (LSTM) ideas (Hochreiter, 1991; Hochreiter & Schmidhuber, 1997b,a), i.e., the constant error carousel and gating, were introduced to overcome the vanishing gradient problem of recurrent neural networks (Hochreiter, 1991; Hochreiter et al., 2000):

The constant error carousel is the additive update of the cell state c t -1 (green) by cell inputs z t and moderated by sigmoid gates (blue). The input gate i t and the forget gate f t control this update, while the output gate o t controls the output of the memory cell, i.e. the hidden state h t . The cell state is normalized or squashed by ψ and then output gating gives the hidden state.

LSTMs have been successfully applied to various domains (Hochreiter et al., 2001, 2007; Schmidhuber, 2015), and prevailed over text generation until the dawn of Transformers in 2017 (Vaswani et al., 2017). The effectiveness of LSTMs has been demonstrated at numerous sequence-related tasks such as generating text (Graves, 2013; Karpathy, 2015), generating handwritings (Graves, 2013), sequence-to-sequence translation (Sutskever et al., 2014), evaluating computer programs (Zaremba & Sutskever, 2014), generating image captions (Karpathy & Fei-Fei, 2015; Hossain et al., 2019), generating source code (Karpathy, 2015), rainfall-runoff modeling (Kratzert et al., 2018, 2019),

or hydrological models for flooding warnings (Nearing et al., 2024). In reinforcement learning, LSTMs are the best performing sequence models, e.g., the AlphaStar model for StarCraft II (Vinyals et al., 2017), the OpenAI Five model for Dota 2 (Karpathy, 2019), and models of the magnetic controller for nuclear fusion (Degrave et al., 2022). LSTMs excel at learning abstractions, i.e., adeptly extracting semantic information and storing it in their memory cells (Karpathy, 2015), which for example became evident by number and syntax neurons (Lakretz et al., 2019), linguistic neurons (Bau et al., 2019), and sentiment neurons (Radford et al., 2017). LSTMs are still used in highly relevant applications (Degrave et al., 2022; Nearing et al., 2024) and have stood the test of time.

Despite their tremendous successes, LSTMs have three main limitations: (i) Inability to revise storage decisions. We exemplify this limitation via the Nearest Neighbor Search problem (see also Appendix D): With a reference vector given, a sequence must be scanned sequentially for the most similar vector in order to provide its attached value at sequence end. The left panel of Figure 1 shows the mean squared error at this task. LSTM struggles to revise a stored value when a more similar vector is found, while our new xLSTM remediates this limitation by exponential gating. (ii) Limited storage capacities, i.e., information must be compressed into scalar cell states. We exemplify this limitation via Rare Token Prediction . In the right panel of Figure 1, the perplexity of token prediction on Wikitext103 (Merity et al., 2017) is given for

partitions of different token frequency. LSTM performs worse on rare tokens because of its limited storage capacities. Our new xLSTM solves this problem by a matrix memory. (iii) Lack of parallelizability due to memory mixing, i.e., the hidden-hidden connections between hidden states from one time step to the next, which enforce sequential processing.

These limitations of LSTM have paved the way for the emergence of Transformers (Vaswani et al., 2017) in language modeling. What performances can we achieve in language modeling when overcoming these limitations and scaling LSTMs to the size of current Large Language Models?

2 Extended Long Short-Term Memory

To overcome the LSTM limitations, Extended Long Short-Term Memory (xLSTM) introduces two main modifications to the LSTM idea of Equation (1). Those modifications - exponential gating and novel memory structures - enrich the LSTM family by two members: (i) the new sLSTM (see Section 2.2) with a scalar memory, a scalar update, and memory mixing, and (ii) the new mLSTM (see Section 2.3) with a matrix memory and a covariance (outer product) update rule, which is fully parallelizable. Both sLSTM and mLSTM enhance the LSTM through exponential gating. To enable parallelization, the mLSTM abandons memory mixing, i.e., the hidden-hidden recurrent connections. Both mLSTM and sLSTM can be extended to multiple memory cells, where sLSTM features memory mixing across cells. Further, the sLSTM can have multiple heads without memory mixing across the heads, but only memory mixing across cells within each head. This introduction of heads for sLSTM together with exponential gating establishes a new way of memory mixing. For mLSTM multiple heads and multiple cells are equivalent.

Integrating these new LSTM variants into residual block modules results in xLSTM blocks (see Section 2.4). Residually stacking those xLSTM blocks in architectures provides xLSTM architectures (see Section 2.4). See Appendix Figure 6 for the xLSTM architecture with its components.

2.1 Review of the Long Short-Term Memory

The original LSTM idea (Hochreiter, 1991; Hochreiter & Schmidhuber, 1997b,a) introduced the scalar memory cell as a central processing and storage unit that avoids vanishing gradients (Hochreiter, 1991; Hochreiter et al., 2000) through the constant error carousel (cell state update). The memory cell contains three gates: input, output, and forget gate. The latter was introduced by Gers et al. (2000). The LSTM memory cell update rules at time step t are:

The weight vectors w z , w i , w f , and w o correspond to the input weight vectors between inputs x t and cell input, input gate, forget gate, and output gate, respectively. The weights r z , r i , r f , and r o correspond to the recurrent weights between hidden state h t -1 and cell input, input gate, forget gate, and output gate, respectively. b z , b i , b f , and b o are the corresponding bias terms. φ and ψ are the cell input and hidden state activation functions (typically tanh ). ψ is used to normalize or squash the cell state, which would be unbounded otherwise. All gate activation functions are sigmoid, i.e., σ ( x ) = 1 / (1+exp( -x )) . In later formulations, multiple scalar memory cells c t ∈ R were combined in a vector c t ∈ R d , which allows the usage of recurrent weight matrices R ∈ R d × d for each gate to mix the cell outputs of memory cells (Greff et al., 2015), for more details see Appendix B.1. Ablation studies showed that all components of the memory cell are crucial (Greff et al., 2015).

2.2 sLSTM

To empower LSTMs with the ability to revise storage decisions, we introduce exponential gates (red) together with normalization and stabilization. In particular, input and forget gates can have exponential activation functions. For normalization, we introduce a normalizer state that sums up the product of the input gate times all future forget gates. The scalar sLSTM forward pass is:

We transfer the original LSTM gating techniques, i.e., input- and/or hidden-dependent gating plus bias term, to the new architectures. Exponential activation functions can lead to large values that cause overflows. Therefore, we stabilize gates with an additional state m t (Milakov & Gimelshein, 2018), see Equations (49) - (51) in the appendix.

New Memory Mixing. sLSTM can have multiple memory cells like the original LSTM (see Appendix B.2). Multiple memory cells enable memory mixing via recurrent connections R z , R i , R f , R o from hidden state vector h to memory cell input z and the gates i , f , o , respectively. A new aspect in memory mixing is the effect of exponential gating. The new sLSTM can have multiple

heads with memory mixing within each head but not across heads. The introduction of heads for sLSTM together with exponential gating establishes a new way of memory mixing.

2.3 mLSTM

To enhance storage capacities of LSTMs, we increase the LSTM memory cell from a scalar c ∈ R to a matrix C ∈ R d × d . Hence, retrieval is performed via a matrix multiplication. At time t , we want to store a pair of vectors, the key k t ∈ R d and the value v t ∈ R d (we use the Transformer terminology). Later at time t + τ , the value v t should be retrieved by a query vector q t + τ ∈ R d . This is the setting of Bidirectional Associative Memories (BAMs) (Kohonen, 1972; Anderson, 1972; Nakano, 1972; Anderson et al., 1977). The covariance update rule (Sejnowski, 1977; Dayan & Willshaw, 1991) for storing a key-value pair is

We assume a layer-norm before projecting inputs to keys and values, therefore they have zero mean. The covariance update rule is optimal (Dayan & Willshaw, 1991) for a maximal separability of retrieved binary vectors, which is equivalent to a maximal signal/noise ratio. Higher separability is possible when limiting retrieval to pairwise interactions and conceding quadratic complexity like attention (Krotov & Hopfield, 2016, 2017; Ramsauer et al., 2021). The covariance update rule is equivalent to Fast Weight Programmers (Schmidhuber, 1992; Schlag et al., 2021), which have later been equipped with a constant decay rate multiplied to C t -1 and a constant learning rate multiplied to v t k ⊤ t (Ba et al., 2016a). In this spirit, we integrate the covariance update rule into the LSTM framework, where the forget gate corresponds to decay rate and the input gate to the learning rate, while the output gate scales the retrieved vector.

For this matrix memory, the normalizer state is the weighted sum of key vectors, where each key vector is weighted by the input gate and all future forget gates. Again, the normalizer state keeps record of the strength of the gates. Since the dot product between query and normalizer state can be close to zero, we use the absolute value of this dot product and lower bound it by a threshold (typically 1.0) as done previously (Sun et al., 2023). The mLSTM forward pass is:

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 49). Since the mLSTM has no memory mixing, this recurrence can be reformulated in a parallel version. For more details we refer to Appendix B.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 Appendix Figure 7 for details. (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. See Appendix Figure 8 for more details. For an xLSTM block containing an sLSTM, we mostly 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.

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 preLayerNorm (Ba et al., 2016b) residual backbones as used in contemporary Large Language Models. See last two columns (from the left) in Figure 6.

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 matrix memory and d × 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

Conceptually, the closest models to xLSTM are Retention (Sun et al., 2023), RWKV (Peng et al., 2023, 2024), GLA (Yang et al., 2023), HGRN2 (Qin et al., 2024) and Mamba (Gu & Dao, 2024). These models share the concepts matrix memory and/or gating. However, in contrast to the new sLSTM, they 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). Other closely related work is discussed further in Appendix C.

4 Experiments

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 have been trained on 15B tokens from SlimPajama (Soboleva et al., 2023). On the same dataset, we perform ablation studies for xLSTM. Then, we compare xLSTM and the best performing methods from Section 4.2 after being trained on 300B tokens from SlimPajama (Soboleva et al., 2023) on downstream tasks, assess their scaling behavior analogous to Kaplan et al. (2020) and Brown et al. (2020), and compare the text generation times and the maximal throughput of the xLSTM in Section 4.3. For all experiments, we use the notation xLSTM[ a : b ] for the ratio 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.

4.1 Synthetic Tasks and Long Range Arena

Firstly, 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 D.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 9. 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. 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 enlarge the context length up to 2048, obtaining extended 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 2. 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 of Appendix D.1.2, indicate that xLSTM's enhanced memory capacities also allow for extrapolating to contexts that are longer than those seen during training.

32

64

128

256

512

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

KV Pairs = 48

32

64

128

256

512

KV Pairs = 96

32

64

128

256

512

KV Pairs = 256

Llama

Mamba

RWKV-5

RWKV-6

xLSTM[1:0]

xLSTM[1:1]

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 D.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. We train models on 15B tokens from SlimPajama (Soboleva et al., 2023), and evaluate their perplexity on the validation set. We compare the following methods: xLSTM, 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), 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]. The models were trained with mixed precision, for RWKV-5, RWKV-6, GLA, HGRN2, the mixedprecision training did not utilize the PyTorch automated mixed precision (see Appendix Section D.2). We categorize the methods into (a) Transformers, (b) State Space Models (SSMs), and (c) Recurrent Neural Networks (RNNs) together with linear Transformers, i.e., linear methods that substitute the Transformer's 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 D.2. Figure 15 in the appendix shows the scaling behaviour for this experiment, indicating that xLSTM will also perform favorably for larger models.

Ablation Studies. Table 1 and Figure 15 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 Appendix Table 6 (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 Appendix Table 6 (bottom), we conclude that having each gate learnable and influenced by the input has an incrementally positive effect. Additional studies on the individual backbone components are discussed in Appendix D.2.

4.3 xLSTM as Large Language Model

Next, we increase the amount of training data to 300B tokens from SlimPajama, the same number of tokens as 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 (Appendix Section D.2) has been found only after the training start of the 300B token experiments (Peng et al., 2024). 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, on downstream tasks, on 471 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 3 for the results. In contrast to other methods, xLSTM models maintain low perplexities for longer contexts.

Model

SlimPajama

(300B) ppl

↓

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 2 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 2 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 D.3.

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. Appendix Table 8 shows token prediction perplexity grouped into language modeling (first seven columns) and fine-grained domain benchmarks (last 5 columns). 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, see Appendix D.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 4 presents the scaling behavior over the number of model parameters. 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. In Figure 16 in Appendix D.3 we plot the scaling behavior over the number of training FLOPs for the Llama baseline and both xLSTM variants. For the xLSTM variants we calculate the FLOPs for the recurrent (see Section 2 and parallel (see Appendix B.3) formulation. The scaling behavior indicates that for larger models xLSTM will continue to perform favourable compared to Transformers and State-Space models.

0

.

2

0

.

4

1

.

0

1

.

4

Number of Parameters

10

9

9

10

11

12

13

14

15

16

17

Validation Perplexity

Llama

Mamba

RWKV-4

xLSTM[7:1]

xLSTM[1:0]

300B Tokens

×

Generation Times and Maximal Throughput. Finally, we measure the text generation time in Figure 5 (left) and the maximal throughput in Figure 5 (right) 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 5 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 5 (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 thus prevents a fast parallel implementation. However, we developed a fast CUDA kernel for sLSTM, which is currently less than two times slower than the parallel mLSTM implementation. (ii) The mLSTM CUDA kernels are not optimized, and therefore the current implementation is about four 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 matrices must be processed. Still, the memory update and retrieval is parameter-free, i.e., parallelizable when using standard matrix operations. Thus, the wall clock time overhead due to the complex memory is minor. (iv) The forget gate initialization must be chosen carefully. (v) Since the matrix memory is sequence length independent, 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 LLM 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

We have partly answered our simple question: How far do we get in language modeling when scaling LSTM to billions of parameters? So far, we can answer: 'At least as far as current technologies like Transformers or State Space Models'. We have enhanced LSTM to xLSTM by exponential gating with memory mixing and a new memory structure. xLSTM models perform favorably on language modeling when compared to state-of-the-art methods like Transformers and State Space Models. The scaling laws indicate that larger xLSTM models will be serious competitors to current LLMs that are built with the Transformer technology. xLSTM has the potential to considerably impact other fields like Reinforcement Learning, Time Series Prediction, or the modeling of physical systems.

References

Achiam, S. Adler, S. Agarwal, et al. GPT-4 technical report. ArXiv , 2303.08774, 2023.

Anderson, J. Silverstein, S. Ritz, and R. Jones. Distinctive features, categorical perception, and probability learning: Some applications of a neural model. Psychological Review , 84:413-451, 1977. doi: 10.1037/0033-295X.84.5.413.

A. Anderson. A simple neural network generating an interactive memory. Mathematical Biosciences , 14, 1972. doi: 10.1016/0025-5564(72)90075-2.

Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré. Zoology: Measuring and improving recall in efficient language models. ArXiv , 2312.04927, 2023.

Ba, G. E. Hinton, V. Mnih, J. Z. Leibo, and C. Ionescu. Using fast weights to attend to the recent past. In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems 29 , pp. 4331-4339. Curran Associates, Inc., 2016a.

Ba, J. R. Kiros, and G. Hinton. Layer normalization. ArXiv , 1607.06450, 2016b.

Bau, Y. Belinkov, H. Sajjad, N. Durrani, F. Dalvi, and J. Glass. Identifying and controlling important neurons in neural machine translation. In International Conference on Learning Representations (ICLR) , 2019. URL https://openreview.net/forum?id=H1z-PsR5KX .

Bisk, R. Zellers, R. LeBras, J. Gao, and Y. Choi. Piqa: Reasoning about physical commonsense in natural language. In AAAI Conference on Artificial Intelligence , volume 34, pp. 7432-7439, 2020.

L. Blodgett, L. Green, and B. O'Connor. Demographic dialectal variation in social media: A case study of African-American English. In Conference on Empirical Methods in Natural Language Processing , pp. 1119-1130, 2016. doi: 10.18653/v1/D16-1120.

Brown, B. Mann, N. Ryder, et al. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems , volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.

M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlós, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, D. B. Belanger, L. J. Colwell, and A. Weller. Rethinking attention with performers. In 9th International Conference on Learning Representations (ICLR) . OpenReview.net, 2021. URL https://openreview.net/forum?id=Ua6zuk0WRH .

Chowdhery, S. Narang, J. Devlin, et al. PaLM: scaling language modeling with pathways. ArXiv , 2204.02311, 2022.

Chronopoulou, M. Peters, and J. Dodge. Efficient hierarchical domain adaptation for pretrained language models. In Conference of the North American Chapter of the Association for Computational Linguistics , pp. 1336-1351, 2022. doi: 10.18653/v1/2022.naacl-main.96.

Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? Try ARC, the AI2 reasoning challenge. ArXiv , 1803.05457, 2018.

M. Cover. Geometrical and statistical properties of systems of linear inequalities with applications in pattern recognition. Electronic Computers, IEEE Transactions on , EC-14(3):326-334, 1965.

Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR) , volume 12, 2024. URL https: //openreview.net/forum?id=mZn2Xyh9Ec .

Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (eds.), Advances in Neural Information Processing Systems (NeurIPS) , 2022. URL https://openreview.net/ forum?id=H4DqfPSibmx .

Dayan and D. J. Willshaw. Optimising synaptic learning rules in linear associative memories. Biological Cybernetics , 65, 1991. doi: 10.1007/bf00206223.

De, S. L. Smith, A. Fernando, A. Botev, G. Cristian-Muraru, A. Gu, R. Haroun, L. Berrada, Y. Chen, S. Srinivasan, G. Desjardins, A. Doucet, D. Budden, Y. W. Teh, R. Pascanu, N. DeFreitas, and C. Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models. ArXiv , 2402.19427, 2024.

Degrave, F. Felici, J. Buchli, et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature , 602:414-419, 2022. doi: 10.1038/s41586-021-04301-9.

Delétang, A. Ruoss, J. Grau-Moya, T. Genewein, L. K. Wenliang, E. Catt, C. Cundy, M. Hutter, S. Legg, J. Veness, and P. A. Ortega. Neural networks and the Chomsky hierarchy. In International Conference on Learning Representations (ICLR) , volume 11, 2023. URL https://openreview. net/forum?id=WbxHAzkeQcn .

Du, Y. Huang, A. M. Dai, et al. GLaM: efficient scaling of language models with mixture-ofexperts. ArXiv , 2112.06905, 2021.

Y. Fu, T. Dao, K. K. Saab, A. W. Thomas, A. Rudra, and C. Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=COZDy0WYGg .

Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy. The Pile: An 800gb dataset of diverse text for language modeling. ArXiv , 2101.00027, 2021.

A. Gers, J. Schmidhuber, and F. Cummins. Learning to forget: Continual prediction with LSTM. Neural Compututation , 12(10):2451-2471, 2000.

Gemini Team Google. Gemini: A family of highly capable multimodal models. ArXiv , 2312.11805, 2023.

Graves. Generating sequences with recurrent neural networks. ArXiv , 1308.0850, 2013.

Greenbaum and G. Nelson. The international corpus of English (ICE) project. World Englishes , 15(1):3-15, 1996.

Greff, R. K. Srivastava, J. Koutník, B. R. Steunebrink, and J. Schmidhuber. LSTM: A search space odyssey. ArXiv , 1503.04069, 2015.

Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. ArXiv , 2312.00752, 2023.

Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. In International Conference on Learning Representations (ICLR) , 2024. URL https://openreview.net/ forum?id=AL1fq05o7H .

Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces. ArXiv , 2111.00396, 2021.

Gupta, A. Gu, and J. Berant. Diagonal state spaces are as effective as structured state spaces. ArXiv , 2203.14343, 2022.

He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pp. 770-778, 2016.

Hochreiter. Untersuchungen zu dynamischen neuronalen Netzen. Master's thesis, Technische Universität München, 1991.

Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation , 9(8):1735-1780, 1997a.

Hochreiter and J. Schmidhuber. LSTM can solve hard long time lag problems. In M. C. Mozer, M. I. Jordan, and T. Petsche (eds.), Advances in Neural Information Processing Systems (NeurIPS) , volume 9, pp. 473-479. MIT Press, Cambridge MA, 1997b.

Hochreiter, Y. Bengio, P. Frasconi, and J. Schmidhuber. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies. In J. Kolen and S. Kremer (eds.), A Field Guide to Dynamical Recurrent Networks . IEEE, 2000.

Hochreiter, A. Steven Younger, and Peter R. Conwell. Learning to learn using gradient descent. In G. Dorffner, H. Bischof, and K. Hornik (eds.), Proc. Int. Conf. on Artificial Neural Networks (ICANN 2001) , pp. 87-94. Springer, 2001.

Hochreiter, M. Heusel, and K. Obermayer. Fast model-based protein homology detection without alignment. Bioinformatics , 23(14):1728-1736, 2007.

Hoffmann, S. Borgeaud, A. Mensch, et al. Training compute-optimal large language models. ArXiv , 2203.15556, 2022.

D. Hossain, F. Sohel, M. F. Shiratuddin, and H. Laga. A comprehensive survey of deep learning for image captioning. ACM Computing Surveys (CSUR) , 51(6):118, 2019.

Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. ArXiv , 2001.08361, 2020.

Karpathy. The unreasonable effectiveness of recurrent neural networks. http://karpathy.github.io/2015/05/21/rnn-effectiveness/, 2015.

Karpathy. OpenAI Five defeats Dota 2 world champions. https://openai.com/research/openai-fivedefeats-dota-2-world-champions, 2019.

Karpathy and L. Fei-Fei. Deep visual-semantic alignments for generating image descriptions. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pp. 3128-3137, 2015.

Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In E. H. Daumé III and A. Singh (eds.), International Conference on Machine Learning (ICML) , volume 119 of Proceedings of Machine Learning Research , pp. 5156-5165. PMLR, 2020.

Katsch. GateLoop: Fully data-controlled linear recurrence for sequence modeling. ArXiv , 2311.01927, 2023.

Kocetkov, R. Li, L. BenAllal, J. Li, C. Mou, C. Mu nozFerrandis, Y. Jernite, M. Mitchell, S. Hughes, T. Wolf, D. Bahdanau, L. vonWerra, and H. deVries. The Stack: 3 TB of permissively licensed source code. ArXiv , 2211.15533, 2022.

Kohonen. Correlation matrix memories. IEEE Transactions on Computers , C-21(4), 1972. doi: 10.1109/tc.1972.5008975.

Kratzert, D. Klotz, C. Brenner, K. Schulz, and M. Herrnegger. Rainfall-runoff modelling using long short-term memory (LSTM) networks. Hydrology and Earth System Sciences , 22(11):6005-6022, 2018.

Kratzert, D. Klotz, G. Shalev, G. Klambauer, S. Hochreiter, and G. Nearing. Benchmarking a catchment-aware long short-term memory network (LSTM) for large-scale hydrological modeling. ArXiv , 1907.08456, 2019.

Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, Deptartment of Computer Science, University of Toronto, 2009.

Krotov and J. J. Hopfield. Dense associative memory for pattern recognition. In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems , pp. 1172-1180. Curran Associates, Inc., 2016.

Krotov and J. J. Hopfield. Dense associative memory is robust to adversarial inputs. ArXiv , 1701.00939, 2017.

Lakretz, G. Kruszewski, T. Desbordes, D. Hupkes, S. Dehaene, and M. Baroni. The emergence of number and syntax units in LSTM language models. In J. Burstein, C. Doran, and T. Solorio (eds.), Conference of the North American Chapter of the Association for Computational Linguistics , pp. 11-20. Association for Computational Linguistics, 2019. doi: 10.18653/v1/N19-1002.

Li, T. Cai, Y. Zhang, D. Chen, and D. Dey. What makes convolutional models great on long sequence modeling? ArXiv , 2210.09298, 2022.

Liang, R. Bommasani, T. Lee, et al. Holistic evaluation of language models. Annals of the New York Academy of Sciences , 1525:140-146, 2023.

Lin, R. Men, A. Yang, C. Zhou, M. Ding, Y. Zhang, P. Wang, A. Wang, L. Jiang, X. Jia, J. Zhang, J. Zhang, X. Zou, Z. Li, X. Deng, J. Liu, J. Xue, H. Zhou, J. Ma, j. Yu, Y. Li, W. Lin, J. Zhou, J. Tang, and H. Yang. M6: A Chinese multimodal pretrainer. ArXiv , 2103.00823, 2021.

Linsley, J. Kim, V. Veerabadran, C. Windolf, and T. Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. Advances in Neural Information Processing Systems (NeurIPS) , 31, 2018.

Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR) , 2019. URL https://openreview.net/forum?id= Bkg6RiCqY7 .

Ma, C. Zhou, X. Kong, J. He, L. Gui, G. Neubig, J. May, and L. Zettlemoyer. Mega: Moving average equipped gated attention. ArXiv , 2209.10655, 2022.

L. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts. Learning word vectors for sentiment analysis. In Annual Meeting of the Association for Computational Linguistics , volume 49, pp. 142-150, 2011.

Magnusson, A. Bhagia, V. Hofmann, et al. Paloma: A benchmark for evaluating language model fit. ArXiv , 2312.10523, 2023.

Mehta, A. Gupta, A. Cutkosky, and B. Neyshabur. Long range language modeling via gated state spaces. ArXiv , 2206.13947, 2022.

Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models. In International Conference on Learning Representations (ICRL) , 2017. URL https://openreview.net/ forum?id=Byj72udxe .

Merrill and A. Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics , 11:531-545, 2023. doi: 10.1162/ tacl_a_00562.

Merrill, J. Petty, and A. Sabharwal. The illusion of state in state-space models. ArXiv , 2404.08819, 2024.

Milakov and N. Gimelshein. Online normalizer calculation for softmax. ArXiv , 1805.02867, 2018.

Nakano. Associatron - a model of associative memory. IEEE Transactions on Systems, Man, and Cybernetics , SMC-2(3):380-388, 1972. doi: 10.1109/TSMC.1972.4309133.

Nearing, D. Cohen, V. Dube, M. Gauch, O. Gilon, S. Harrigan, A. Hassidim, D. Klotz, F. Kratzert, A. Metzger, S. Nevo, F. Pappenberger, C. Prudhomme, G. Shalev, S. Shenzis, T. Y. Tekalign, D. Weitzner, and Y. M. B. Kosko. Global prediction of extreme floods in ungauged watersheds. Nature , 627:559-563, 2024. doi: 10.1038/s41586-024-07145-1.

Olsson, N. Elhage, N. Nanda, et al. In-context learning and induction heads. ArXiv , 2209.11895, 2022.

Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. In Proceedings of the 40th International Conference on Machine Learning (ICML) . JMLR.org, 2023. doi: 10.5555/3618408.3619518.

Papasavva, S. Zannettou, E. DeCristofaro, G. Stringhini, and J. Blackburn. Raiders of the lost KeK: 3.5 years of augmented 4chan posts from the politically incorrect board. In International AAAI Conference on Web and Social Media (ICWSM) , volume 14, pp. 885-894, 2020.

Paperno, G. Kruszewski, A. Lazaridou, N.-Q. Pham, R. Bernardi, S. Pezzelle, M. Baroni, Gemma G. Boleda, and R. Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Annual Meeting of the Association for Computational Linguistics , volume 1, pp. 1525-1534, 2016.

Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, A. Cappelli, H. Alobeidli, B. Pannier, E. Almazrouei, and J. Launay. The RefinedWeb dataset for Falcon LLM: Outperforming curated corpora with web data, and web data only. ArXiv , 2306.01116, 2023.

Peng, E. Alcaide, Q. Anthony, et al. RWKV: Reinventing RNNs for the transformer era. ArXiv , 2305.13048, 2023.

Peng, D. Goldstein, Q. Anthony, et al. Eagle and Finch: RWKV with matrix-valued states and dynamic recurrence. ArXiv , 2404.05892, 2024.

Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré. Hyena hierarchy: Towards larger convolutional language models. In Proceedings of the 40th International Conference on Machine Learning (ICML) . JMLR.org, 2023. doi: 10.5555/3618408.3619572.

Poli, A. W. Thomas, E. Nguyen, P. Ponnusamy, B. Deiseroth, K. Kersting, T. Suzuki, B. Hie, S. Ermon, C. Ré, C. Zhang, and S. Massaroli. Mechanistic design and scaling of hybrid architectures. ArXiv , 2403.17844, 2024.

Qin, S. Yang, and Y. Zhong. Hierarchically gated recurrent neural network for sequence modeling. In Advances in Neural Information Processing Systems (NeurIPS) , volume 37, 2023. URL https://openreview.net/forum?id=P1TCHxJwLB .

Qin, S. Yang, W. Sun, X. Shen, D. Li, W. Sun, and Y. Zhong. HGRN2: Gated linear RNNs with state expansion. ArXiv , 2404.07904, 2024.

R. Radev, P. Muthukrishnan, and V. Qazvinian. The ACL anthology network corpus. In Workshop on Text and Citation Analysis for Scholarly Digital Libraries (NLPIR4DL) , pp. 54-61. Association for Computational Linguistics, 2009.

Radford, R. Jozefowicz, and I. Sutskever. Learning to generate reviews and discovering sentiment. ArXiv , 1704.01444, 2017.

Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. https://openai.com/index/better-language-models , 2019.

W. Rae, S. Borgeaud, T. Cai, et al. Scaling language models: Methods, analysis & insights from training Gopher. ArXiv , 2112.11446, 2021.

Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. ArXiv , 1910.10683, 2019.

Ramsauer, B. Schäfl, J. Lehner, P. Seidl, M. Widrich, L. Gruber, M. Holzleitner, M. Pavlovi´ c, G. K. Sandve, V. Greiff, D. Kreil, M. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations (ICLR) . OpenReview, 2021.

Reid, V. Zhong, S. Gururangan, and L. Zettlemoyer. M2D2: A massively multi-domain language modeling dataset. In Conference on Empirical Methods in Natural Language Processing , pp. 964-975, 2022.

Reid, N. Savinov, D. Teplyashin, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. ArXiv , 2403.05530, 2024.

H. Ribeiro, J. Blackburn, B. Bradlyn, E. DeCristofaro, G. Stringhini, S. Long, S. Greenberg, and S. Zannettou. The evolution of the manosphere across the web. In Proceedings of the international AAAI conference on web and social media , volume 15, pp. 196-207, 2021.

Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM , 64(9):99-106, 2021.

L. Scao, A. Fan, C. Akiki, et al. BLOOM: A 176B-parameter open-access multilingual language model. ArXiv , 2211.05100, 2022.

Schlag, K. Irie, and J. Schmidhuber. Linear transformers are secretly fast weight programmers. In M. Meila and T. Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning (ICML) , volume 139 of Proceedings of Machine Learning Research , pp. 9355-9366. PMLR, 2021.

Schmidhuber. Learning to control fast-weight memories: An alternative to recurrent nets. Neural Computation , 4(1):131-139, 1992.

Schmidhuber. Deep learning in neural networks: An overview. Neural Networks , 61:85-117, 2015. doi: 10.1016/j.neunet.2014.09.003.

Schulman, B. Zoph, C. Kim, J. Hilton, et al. ChatGPT: Optimizing language models for dialogue. https://openai.com/blog/chatgpt/, 2022. OpenAI Research.

J. Sejnowski. Storing covariance with nonlinearly interacting neurons. Journal of Mathematical Biology , 4, 1977. doi: 10.1007/BF00275079.

Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. ArXiv , 1909.08053, 2019.

T. H. Smith, A. Warrington, and S. W. Linderman. Simplified state space layers for sequence modeling. ArXiv , 2208.04933, 2022.

Soboleva, F. Al-Khateeb, R. Myers, J. R. Steeves, J. Hestness, and N. Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. https://www.cerebras.net/blog/ slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama , 2023. URL https://huggingface.co/datasets/cerebras/SlimPajama-627B .

Soldaini, R. Kinney, A. Bhagia, et al. Dolma: an open corpus of three trillion tokens for language model pretraining research. ArXiv , 2306.01116, 2023.

Soltan, S. Ananthakrishnan, J. FitzGerald, R. Gupta, W. Hamza, H. Khan, C. Peris, S. Rawls, A. Rosenbaum, A. Rumshisky, C. S. Prakash, M. Sridhar, F. Triefenbach, A. Verma, G. Tur, and P. Natarajan. AlexaTM 20B: Few-shot learning using a large-scale multilingual Seq2Seq model. ArXiv , 2208.01448, 2022.

K. Srivastava, K. Greff, and J. Schmidhuber. Training very deep networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems (NeurIPS) , volume 28. Curran Associates, Inc., 2015.

Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei. Retentive network: A successor to transformer for large language models. ArXiv , 2307.08621, 2023.

Sutawika, L. Gao, H. Schoelkopf, et al. EleutherAI/lm-evaluation-harness: Major refactor, 2023.

Sutskever, O. Vinyals, and Q. V. V. Le. Sequence to sequence learning with neural networks. In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 27 (NIPS'13) , pp. 3104-3112. Curran Associates, Inc., 2014.

Tay, D. Bahri, D. Metzler, D.-C. Juan, Z. Zhao, and C. Zheng. Synthesizer: Rethinking selfattention in transformer models. ArXiv , 2005.00743, 2020.

Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations (ICRL) , 2021. URL https://openreview.net/forum?id=qVyeW-grC2k .

Thoppilan, D. deFreitas, J. Hall, et al. LaMDA: Language models for dialog applications. ArXiv , 2201.08239, 2022.

TogetherComputer. Redpajama: an open dataset for training large language models, 2023. URL https://github.com/togethercomputer/RedPajama-Data .

Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models. ArXiv , 2302.1397, 2023.

Vadas and J. R. Curran. Parsing noun phrases in the Penn Treebank. Computational Linguistics , 37(4):753-809, 2011.

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS) , volume 30, pp. 5998-6008. Curran Associates, Inc., 2017.

Vinyals, T. Ewalds, S. Bartunov, et al. Starcraft II: A new challenge for reinforcement learning. ArXiv , 1708.04782, 2017.

Wang, J. N. Yan, A. Gu, and A. M. Rush. Pretraining without attention. ArXiv , 2212.10544, 2022.

Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. ArXiv , 2006.04768, 2020.

Wang, Y. Sun, Y. Xiang, et al. ERNIE 3.0 Titan: Exploring larger-scale knowledge enhanced pre-training for language understanding and generation. ArXiv , 2112.12731, 2021.

Wu and K. He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pp. 3-19, 2018.

Xue, N. Constant, A. Roberts, M. Kale, R. Al-Rfou, A. Siddhant, A. Barua, and C. Raffel. mT5: A massively multilingual pre-trained text-to-text transformer. In Conference of the North American Chapter of the Association for Computational Linguistics , pp. 483-498, 2021. doi: 10.18653/v1/2021.naacl-main.41.

Yang and Y. Zhang. FLA: A Triton-based library for hardware-efficient implementations of linear attention mechanism, 2024. URL https://github.com/sustcsonglin/ flash-linear-attention .

Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardwareefficient training. ArXiv , 2312.06635, 2023.

Zannettou, B. Bradlyn, E. DeCristofaro, H. Kwak, M. Sirivianos, G. Stringini, and J. Blackburn. What is Gab: A bastion of free speech or an alt-right echo chamber. In The Web Conference , pp. 1007-1014, 2018. doi: 10.1145/3184558.3191531.

Zaremba and I. Sutskever. Learning to execute. ArXiv , 1410.4615, 2014.

Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. HellaSwag: Can a machine really finish your sentence? In Annual Meeting of the Association for Computational Linguistics , pp. 4791-4800, 2019.

Zeng, X. Liu, Z. Du, et al. GLM-130B: An open bilingual pre-trained model. ArXiv , 2210.02414, 2022.

Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, T. Mihaylov, M. Ott, S. Shleifer, K. Shuster, D. Simig, P. S. Koura, A. Sridhar, T. Wang, and L. Zettlemoyer. OPT: Open pre-trained transformer language models. ArXiv , 2205.01068, 2022.

Contents

A Broader Impacts

Our work introduces novel LSTM architectures that perform favorably compared to Transformers or State Space Models on language modeling. As our novel LSTM models are inherently recurrent and demonstrably extrapolate well to large contexts, they have the potential for near constant scaling at inference. This compares to linear scaling in the context length for Transformer architectures, representing a vast potential in energy and emission savings when deploying such models to real world applications. Thus, in an ever more connected, data-rich world with ubiquitous compute, our work has the potential to increase general accessibility of state-of-the-art machine learning models whilst making them more environmentally sustainable.

However, any novel machine learning technique that has the potential to impact real world applications can be used for harm, not just for good. Our models with their beneficial inference economics and their ability to extrapolate to large contexts at no overhead could be used, for example, to generate and spread disinformation on a grander scale. Moreover, savings in energy and emissions afforded by deploying our novel xLSTM models might be outweighed by an increase in demand for such models, leading to an overall increase in energy consumption and environmentally harmful emissions. Whilst our work focuses on new, more efficient Language Model architectures, it does not directly address the issues of fairness and bias of such models.

B Extended Long Short-Term Memory

Memory Cells

🠆

Constant Error Carousel

🠆

🠆

Recurrent Inference

🠆

Recurrent Training

sLSTM

+ New Memory Mixing

Memory Cells

xLSTM Blocks

xLSTM

mLSTM

+ Exponential Gating

+ Parallel Training

+ Covariance Update Rule

+ Matrix Memory

LSTM

+ Exponential Gating

B.1 Vanilla Long Short-Term Memory Formulation: Vector Notation

The vanilla LSTM memory cell update rules (Greff et al., 2015) at time step t extend the scalar cell state formulation to a vector of cell states:

The matrices W z , W i , W f , and W o correspond to the input weights between inputs x t and cell input, input gate, forget gate, and output gate, respectively. The matrices R z , R i , R f , and R o correspond to the recurrent weights between hidden state h t -1 and cell input, input gate, forget gate, and output gate, respectively. b z , b i , b f , and b o are the corresponding bias vectors. φ and ψ are the cell input and hidden state activation functions (typically tanh ). ψ is used to normalize or squash the cell state, which would be unbounded otherwise.

B.2 sLSTM

Similar to the LSTM in Section B.1, also the sLSTM can be vectorized to multiple cells:

Here, the cell input activation function φ is tanh , the hidden state activation function is the identity. φ helps stabilizing the recurrence.

Considering external gradient contribution δ ext h t from subsequent layers and recurrent gradient contribution δ R h t from gradients from future states flowing over the cell interaction matrix R , we obtain the recursive backward pass of sLSTM, where δ a indicates gradients with respect to parameter / internal variable a :

with the derivatives of the respective gate activation function i ′ t = exp ′ ( ˜ i t ) = exp( ˜ i t ) = i t , o ′ t = σ ′ (˜ o t ) , and f ′ t = σ ′ ( ˜ f t ) or f ′ t = f t depending on the forget gate activation. φ ′ ( z ) is the derivative of the cell input activation function φ ( z ) .

The matrices R z , R i , R f , R o are block-diagonal which is analogous to multiple heads in the mLSTM. This way, the parameters reduce to d 2 / ( N h ) , where N h is the number of heads, limiting the cell interactions to individual heads. This parameter efficient formulation of cell interactions together with the exponential gating is called the new memory mixing. Finally, to stabilize the backward pass, we clip the magnitude of δ R h t to 10 , as a means to prohibit exploding gradients for long context lengths.

sLSTM Stabilized Version. The stabilized version of sLSTM introduces a new stabilizer state m , applied as:

We show that replacing f t by f ′ t and i t by i ′ t in the forward pass does neither change the output of the whole network nor the derivatives of the loss with respect to the parameters.

The stabilization state m has no gradient, and hence does not influence the other gradients. We re-define c ( s ) t and n ( s ) t as stabilized cell and normalizer states:

Inserting Equation 49 into Equation 8 yields:

Therefore, since the loss solely depends on h t , there's no dependency on m t , and consequently, no gradient exists for this stabilization state. Note that m t can be chosen arbitrarily. We choose m t = max(log ( f t ) + m t -1 , log ( i t )) , which stabilizes the exponential function. One can even find m t , such that the normalizer state n t can be eliminated, but this version was experimentally found to be numerically unstable in the backward pass.

B.3 mLSTM

Throughout this section, 1 ∈ R T denotes a column vector of ones and 1 ⊤ ∈ R 1 × T a row vector of ones, where T is the dimension of this vector space.

Recurrent mLSTM Backward Pass. The recurrent formulation of the mLSTM cell in Equation 16 yields the following backward pass recurrence, where δ a indicates gradients with respect to parameter or internal variable a and δ ext h t denotes gradients from subsequent layers:

and Ω( z ) = Θ ( z -1) -Θ( -z -1) , Θ( z ) being the Heaviside step function. γ ( z ) is either σ ′ ( z ) or exp( z ) , depending on the forget gate activation.

Parallel mLSTM Forward Pass. The mLSTM recurrence in Equations (16-24) can be reformulated in a parallel form, which is used to speed up training. After training we can still use the recurrent formulation for fast text generation.

Instead of processing each input x t ∈ R d at time step t sequentially, the parallel version processes all timesteps of a full sequence X ∈ R T × d at once, where T is the sequence length and d is the head dimension. We present the forward pass of the mLSTM for a single head and drop the head dimension for simplicity.

Let ˜ f ∈ R T be the forget gate pre-activations and ˜ i ∈ R T be the input gate pre-activations for a full sequence. We construct the forget gate activation matrix F ∈ R T × T by

and the input gate pre-activation matrix ˜ I ∈ R T × T by

By applying the elementwise exponential input gate activation function naively, we obtain the unstabilized gate activation matrix D ∈ R T × T as

In order to avoid overflow due to the exponential function we apply the same stabilization as in the recurrent sLSTM, see Equation 49. In the parallel formulation of the mLSTM we get a numerically stable gate activation matrix D ′ ∈ R T × T by taking the logarithm of D element-wise and subtracting the row-wise maximum value of D from each element:

Given the queries, keys and values Q , K , V ∈ R T × d , for a full sequence we can compute all hidden pre-activation states ˜ H ∈ R T × d in parallel for the un-stabilized version by

Note that we extract the 1 √ d factor for K explicitly here and further on. For the stabilized version this yields

where for both versions the hidden pre-activation states ˜ H are identical.

With the output gate pre-activations ˜ O ∈ R T × d we can compute the hidden states H ∈ R T × d for all timesteps by applying the output gate in parallel for each timestep element-wise:

This gives the parallel forward pass of the mLSTM for a full input sequence X ∈ R T × d .

Parallel mLSTM Backward Pass. We present the backward pass of the mLSTM for the stabilized version only. For completeness we summarize the forward pass in the stabilized version before we present the backward pass.

Given the forget gate matrix F ∈ R T × T , the logarithm of the forget gate matrix F = log F ∈ R T × T , and the input gate matrix I ∈ R T × T as introduced above, together with the queries, keys and values Q , K , V ∈ R T × d , we can write the forward pass of the mLSTM in the stabilized version as:

With this forward pass we can compute the gradients δ a for all intermediate and input variables to the mLSTM forward pass in the backward pass. We denote the gradient with respect to variable a as δ a .

Given the output gradient δ ˜ H ∈ R T × d we can compute the backward pass for the intermediate gradients as:

We do not compute the gradients for m as they cancel out (see the proof in the recurrent sLSTM).

With these intermediate gradients the gradients for the logarithmic forget gate matrix δ F ∈ R T × T , the input gate matrix δ I ∈ R T × T , and the queries, keys and values δ Q , δ K , δ V ∈ R T × d are given by

Having computed the gradients for the logarithmic forget gate matrix δ F , we can compute the gradients for the forget gate pre-activations δ ˜ f = [ δ ˜ f 1 , δ ˜ f 2 , ..., δ ˜ f T ] ⊤ ∈ R T .

Recall the logarithmic forget gate matrix F = log F is computed by

With the substitution f = log σ ( ˜ f ) we compute the gradients for the logarithmic forget gate activations δ f = [ δ f 1 , δ f 2 , ..., δ f T ] ⊤ ∈ R T as

where the last equation makes use of the following:

Finally, we compute the input gate pre-activations' gradients δ ˜ i = [ δ ˜ i 1 , δ ˜ i 2 , ..., δ ˜ i S ] ⊤ ∈ R T as the column-wise sum over the rows of the input gate matrix δ I :

This completes the backward pass of the parallel mLSTM for a full input sequence X ∈ R T × d .

B.4 Detailed Block Structure

PF=¾

NH=4

NH=4

NH=4

f

z

o

Conv4

Swish

i

NH=4

NH=4

NH=4

NH=4

NH=4

sLSTM

GN

LN

LN

PF=

PF=

3

4

3

4

GeLU

PF=2

PF=½

PF=2

NH=4

BS=4

BS=4

BS=4

Conv4

LSkip

Swish

Swish

q

k

v

i

f

mLSTM

GN

LN

C Extended 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, 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, 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).

D Experiments

Training Setup. For all experiments, we use Python 1 3.11 with PyTorch 2.2.0 2 , and CUDA 12.1 3 . 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] , 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 1 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.

D.1 Synthetic Tasks and Long Range Arena

D.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 based on 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 , β 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.

Missing

Duplicate

Mod

Arithmetic

(w Brackets)

Solve

Equation

Cycle Nav

Even Pairs

Mod

Arithmetic

(w/o Brackets)

Parity

Majority

Majority

Count

Llama

Mamba

Retention

Hyena

RWKV-4

RWKV-5

RWKV-6

LSTM

(Block)

LSTM

xLSTM[0:1]

xLSTM[1:0]

xLSTM[1:1]

0.92

±

0.02

0.08

±

0.0

0.02

±

0.0

0.02

±

0.0

0.04

±

0.01

1.0

±

0.0

0.03

±

0.0

0.03

±

0.01

0.37

±

0.01

0.13

±

0.0

0.69

±

0.0

0.15

±

0.0

0.04

±

0.01

0.05

±

0.02

0.86

±

0.04

1.0

±

0.0

0.05

±

0.02

0.13

±

0.02

0.69

±

0.01

0.45

±

0.03

0.13

±

0.01

0.03

±

0.0

0.03

±

0.0

0.03

±

0.0

0.05

±

0.01

0.51

±

0.07

0.04

±

0.0

0.05

±

0.01

0.36

±

0.0

0.12

±

0.01

0.3

±

0.02

0.06

±

0.02

0.05

±

0.0

0.02

±

0.0

0.06

±

0.01

0.93

±

0.07

0.04

±

0.0

0.04

±

0.0

0.36

±

0.01

0.18

±

0.02

0.54

±

0.0

0.21

±

0.01

0.06

±

0.0

0.07

±

0.0

0.13

±

0.0

1.0

±

0.0

0.07

±

0.0

0.06

±

0.0

0.63

±

0.0

0.13

±

0.0

0.49

±

0.04

0.15

±

0.01

0.08

±

0.0

0.08

±

0.0

0.26

±

0.05

1.0

±

0.0

0.15

±

0.02

0.06

±

0.03

0.73

±

0.01

0.34

±

0.03

0.96

±

0.0

0.23

±

0.06

0.09

±

0.01

0.09

±

0.02

0.31

±

0.14

1.0

±

0.0

0.16

±

0.0

0.22

±

0.12

0.76

±

0.01

0.24

±

0.01

0.99

±

0.0

0.15

±

0.0

0.76

±

0.0

0.5

±

0.05

0.97

±

0.03

1.0

±

0.0

0.91

±

0.09

1.0

±

0.0

0.58

±

0.02

0.27

±

0.0

0.94

±

0.01

0.2

±

0.0

0.72

±

0.04

0.38

±

0.05

0.93

±

0.07

1.0

±

0.0

1.0

±

0.0

1.0

±

0.0

0.82

±

0.02

0.33

±

0.0

0.84

±

0.08

0.23

±

0.01

0.57

±

0.09

0.55

±

0.09

1.0

±

0.0

1.0

±

0.0

1.0

±

0.0

1.0

±

0.0

0.75

±

0.02

0.22

±

0.0

0.97

±

0.0

0.33

±

0.22

0.03

±

0.0

0.03

±

0.01

0.86

±

0.01

1.0

±

0.0

0.04

±

0.0

0.04

±

0.01

0.74

±

0.01

0.46

±

0.0

0.7

±

0.21

0.2

±

0.01

0.15

±

0.06

0.24

±

0.04

0.8

±

0.03

1.0

±

0.0

0.6

±

0.4

1.0

±

0.0

0.64

±

0.04

0.5

±

0.0

Context Sensitive

Deterministic

Context Free

Regular

Additional Formal Language Results. Figure 9 and Figure 10 showcase supplementary results of the formal language tasks. The former extends the results with additional models. The latter details 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 | and the random accuracy 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. 1

Cycle Nav Given a string of 'movement tokens' ( +1 , -1 , STAY) compute the end position of the agent with start position 0. The position must be computed modulo the maximum position.

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.

Majority Given a string of tokens, compute the token that occurred most often in the sequence.

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.

•

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.

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]).

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]).

Odds First An string of tokens t 1 , t 2 , t 3 , ...t n is given. Output all tokens with and odd index ( t 1 , t 3 , ... ) then the token with an even index ( t 2 , t 4 ,..) . Apart from that keep the ordering of the initial string.

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.

Repetition Given a string of tokens - repeat it.

Reverse String Given a string of tokens - repeat it in reverse order.

Stack Manipulation An initial stack content is given, followed by a sequence of push and pop operations. Compute the stack content after the operations

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.

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],

D.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 11 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 12. 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 13. 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 14 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.

D.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]] - → 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).

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 ) is set to 3. The model hyperparameters for each dataset are displayed in Table 3.

Results. Table 4 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.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Transformer

Context Length = 64

Context Length = 128

Context Length = 256

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

(x)LSTM-Family

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Mamba

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

RWKV

32

64

128

256

512

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Others

32

64

128

256

512

32

64

128

256

512

Llama

Mamba

Mamaba (noWT)

RWKV-4

RWKV-5

RWKV-6

Retention

Hyena

H3

xLSTM[0:1]

xLSTM[1:0]

xLSTM[1:1]

LSTM (Block)

LSTM

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Context Length = 756

Key-Value Pairs = 48

Key-Value Pairs = 96

Key-Value Pairs = 256

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Context Length = 1024

32

64

128

256

512

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Context Length = 2048

32

64

128

256

512

32

64

128

256

512

Llama

Mamba

RWKV-5

RWKV-6

xLSTM[1:0]

xLSTM[1:1]

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Context Length = 4096

Key-Value Pairs = 48

Key-Value Pairs = 96

Key-Value Pairs = 256

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Context Length = 6144

128

256

512

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Context Length = 8192

128

256

512

128

256

512

Mamba

RWKV-5

RWKV-6

xLSTM[1:0]

xLSTM[1:1]

D.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) 5 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 , β 2 )=(0.9, 0.95) and an epsilon parameter of 1e-5, and gradient clipping at gradient norm 1. 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. Except for gates we do not use biases in our models, e.g. in other linear layers. We do not apply weight decay to biases and LayerNorm weights. 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 where applicable, 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.

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) 6 , RWKV-5, RWKV-6 (Peng et al., 2024) 7 . 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 by the authors of GLA (Yang et al., 2023; Yang & Zhang, 2024) 8 . 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 models using the flash-linear-attention (Yang & Zhang, 2024) library because of compilation problems and Transformers as for those training curves deviated.

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 9 . 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.

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 did not work because of the custom kernels.

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 8) and the sLSTM in the post up-projection block

0

.

2

0

.

4

1

.

0

1

.

4

2

.

7

Number of Parameters

10

9

10

11

12

13

14

15

16

17

18

19

Validation Perplexity

Llama

Mamba

RWKV-4

xLSTM[7:1]

xLSTM[1:0]

15B Tokens

Figure 15: Method comparison on next token prediction when trained on 15B tokens from SlimPajama. Performance measure in validation perplexity for the best methods of each model class (see Table 1) are reported. The performance degradation of xLSTM[7:1] at 2.7B is due to initially slower training convergence that leads to an especially undertrained model. xLSTM is the best method at all sizes.

×

Table 6: Ablation studies. Top: Ablation studies on the new xLSTM components, contributing the strong performance improvement of xLSTM over vanilla LSTM to both the exponential gating and the matrix memory. Bottom: Ablation studies on different gating techniques. We consider an xLSTM[1:0] with sigmoid forget gate and exponential input gate. Bias initialization ∞ means that the forget gate is set to one, [3 , 6] indicates that values are taken equidistant in the respective interval, and N (0 , 0 . 1) that values are randomly chosen from a Gaussian with mean 0 and std 0 . 1 . PPL denotes validation perplexity. The first two lines correspond to models similar to linearized attention, line four to Retention, line five to RWKV-5, and line six to RWKV-6. Dependencies of the gates on the input lead to better performance.

Ablation studies on different gating techniques.

N

(see Figure 7). In this experiment we study the 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 7) with 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 7 shows that the matrix memory benefits from the pre up-projection block structure, and that the convolution within this block is important.

Details on new xLSTM Components Ablation Study. In Table 6 (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 7), further boosts performance. Adding Exponential Gating to this

architecture yields the sLSTM as depicted in Figure 7, with another large performance improvement. Finally, adding the best Matrix Memory variant found in Table 7 by replacing some sLSTM blocks with the mLSTM (see Figure 8) gives xLSTM[7:1] with the best performance.

Details on Gating Technique Ablation Study. In Table 6 (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 5. For Llama and Mamba we use the learning rates reported by Gu & Dao (2023).

D.3 xLSTM Large Language Models - SlimPajama300B

General Training Procedure. We use the same general training procedure as in Section D.2 with peak learning rates from Table 5. 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.

Additional Scaling Law Plots with Number of FLOPs. In Figure 16 we plot the scaling behavior of our xLSTM variants over the number of training flops and compare to the Llama baseline. We compute the FLOPs for the recurrent (see Section 2) and the parallel (see Section B.3) formulation of the mLSTM. We can see that for the larger models (760M and 1.3B) the recurrent xLSTM FLOP counts are larger than the Llama FLOP counts. The reason for this is that we kept the number of heads constant when increasing the model size (see Table 5) while increasing the head dimension. In contrast Llama increases the number of heads while keeping the head dimension constant.

10

21

Number of FLOPs

9

10

11

12

13

14

15

16

17

Validation Perplexity

Llama

xLSTM[1:0] recurrent

xLSTM[1:0] parallel

xLSTM[7:1] recurrent

xLSTM[7:1] parallel

300B Tokens

Details on FLOP computation. For computing the number of FLOPs we follow the procedure from Hoffmann et al. (2022). We include all training FLOPs, including those contributed by the embedding matrices. We assume that the backward pass has twice the number of FLOPs of the forward pass. For the forward pass, the number of FLOPs of the mLSTM and sLSTM for a single sequence can be approximated by:

mLSTM (recurrent): num_heads × seq_len × (6 × head_dim × head_dim + 8 × head_dim)

sLSTM: num_heads × seq_len × head_dim × (8 × head_dim + 12)

mLSTM (parallel): num_heads × seq_len × seq_len × (4 × head_dim + 8)

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 10 . 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 are shown in Table 8.

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 E.

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 8 as the exponential of the negative log-likelihoods 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)

E 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 8 are not macro averages of the values shown in Table 10.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: See sections 2 and 4.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See Section 5.

Guidelines:

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide one proof in the Appendix, which is clearly linked to the main paper. All assumptions are stated. We enumerate and cross-reference all equations.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We state all experimental details in Appendix D. We state which datasets we used and provide references. We provide the model code.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The data is publicly available. The code is contained in the submission and will be also released publicly.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See Appendix D and attached code.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide error bars only for small scale experiments. For large scale experiments, providing error bars is computationally too expensive.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Section D we state the compute resources used for this project.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We refer to Section A

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [Yes]

Justification: As of now we do release the weights of our pretrained models.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification:

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide the model code with this submission.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We rely on openly available datasets and reference them properly.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the Extended Long Short-Term Memory (xLSTM), which enhances traditional LSTMs with exponential gating and new matrix memory. These improvements address LSTM limitations and enhance their memory storage capability, as well as the ability to revise storage decisions and be more parallelizable. The authors test the efficacy of their method with a set of very exhaustive experiments, which confirm that xLSTMs are competitive with both SSMs and Transformer-based architectures.\n\nSTRENGTHS:\nI believe this paper is exceptionally strong and represents a highly innovative contribution to the field of sequence modeling. It is compelling to see that traditional RNN-based architectures, with some modifications, can outperform Transformer-based architectures, which aligns with some of the recent trends in this field [1,2]. I hope this direction gains traction within the community. This paper also reinforces the notion that old ideas, with slight adjustments, can effectively compete with the most popular models today. Time will tell if these architectures will become a viable alternative to their attention-based counterparts.\n\nHere are some of the key strengths of the paper:\n\n1. An innovative approach to replacing input gates in LSTMs to improve storage decisions.\n2. An intuitive methodology for storing and retrieving information in memory in matrix form.\n3. Linear computation and constant memory complexity with respect to sequence length.\n4. Extremely exhaustive evaluation in both language-based and synthetic tasks.\n\n[1] Gu, Albert, Karan Goel, and Christopher Ré. \"Efficiently modeling long sequences with structured state spaces.\" arXiv preprint arXiv:2111.00396 (2021).\n\n[2] Orvieto, Antonio, et al. \"Resurrecting recurrent neural networks for long sequences.\" International Conference on Machine Learning. PMLR, 2023.\n\nWEAKNESSES:\nI think the paper has been executed almost flawlessly, so I will forgo this section and will simply add some questions in the section below.\n\nQUESTIONS:\nFor my interest, I would like the authors to clarify a few things:\n\n1. Structure of Memory Mixing Matrices: As far as I can tell, the idea of matrix-based updates has been used for a while now. I do appreciate the introduction of different “heads” through a block-diagonal structure. In some sense, this approach could be viewed as restrictive, as it imposes a specific form on the update matrix. I wonder if there might be any benefit to considering a more general structure (which could be viewed as a sort of graph) that could be learned from the downstream task loss? \n\n2. State Normalization: The idea of normalizing the state to avoid overflow using the additional state $m_t$ is clever. However, at some point I pressume $i_t$ has to be calculated and stored as well. Could this still lead to some overflow issues? What did you observe in practice?\n\n3. Matrix Memory: I find this idea very intuitive. While the authors have focused on using this for the cell state of the LSTM in this context, I wondered if they had tested the same idea for other (even vanilla) RNNs. Should similar improvements in memory performance be expected intuitively?\n\n4. Overall Architectural Decisions: What was the thought process behind designing the whole block structure (e.g., 1D convolutional layers, etc.) shown in the appendices? Was this inspired by previous high-performing architectures in sequence modeling (i.e. previous high-performing sequence modeling architectures)? Is this something that the authors tuned significantly, or is there potential for even better performance?\n\n5. Extrapolation: Do the results in Figure 3 suggest that xLSTMs could train on shorter context lengths while retaining the ability to extrapolate to longer sequences? Note that this is just a question to get the authors' opinion; I am not suggesting that they should run experiments along these lines, as the cost would be too high.\n\n6. Other Applications: This paper has focused on text data as well as other synthetic tasks, but I was wondering if the authors would expect it to work well for other data types (e.g., temporal or image data)?\n\nLIMITATIONS:\nYes, the limitations of the proposed method are mentioned in the paper in a dedicated section.", "TITLE": "Official Review", "RECOMMENDATION": 10, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nProposes some extension to LSTM. Specifically:\n\n* exponential gating with appropriate normalization and stabilization techniques\n\n* sLSTM with a scalar memory, a scalar update, and new memory mixing. This still needs to be calculated sequentially.\n\n* mLSTM that is fully parallelizable with a matrix memory (instead of just a vector) and a covariance update rule\n\n* putting sLSTM & mLSTM together into a residual block: the xLSTM block\n\n* Extended LSTM (xLSTM) blocks are residually stacked into xLSTM architecture\n\nTested for language modeling.\n\nCompared to Transformers and State Space Models.\n\nSTRENGTHS:\nScales up LSTM models.\n\nProposes extensions to the standard LSTM model which perform really well.\n\nLots of experiments. (Although most is only in the appendix...)\n\nOpen code.\n\nWEAKNESSES:\nSome of the notation is confusing. Specifically, I'm never really sure whether it's about scalars or vectors. I also think there are some errors regarding this. I don't understand why to write it at all in scalar notation.\n\nThere is never really an exact definition of the whole model. We can only infer it from the figures. (Or the code.)\n\nThe scaling laws comparison is a bit flawed: I don't think you can just compare the number of parameters across different architectures. More reasonable would be to compare inference time.\n\nQUESTIONS:\nIn all the formulas, it's not immediately clear whether e.g. h_t, c_t etc are scalar values or vectors. I assume those are vectors. It would be good to write that explicitly (sth like c_t, h_t, ... \\in \\R^D) or so.\n\nHowever, then it says that w_z, w_i, w_f, w_o are also vectors (eq 4-6). This is weird. That should be matrices, or not? Also the r_z, r_i, r_f, r_o, they are just called \"weights\" (p.3 line 94), but that should be matrices as well, specifically all in \\R^{D \\times D}. It would help to state that explicitly. The notation is also a bit uncommon. Most common is to use capital letters for matrices.\n\nOr is h_t here really a scalar, not a vector, e.g. in eq 2? But then, in eq 4-7, the multiplication with r is wrong? Or are those r also scalars? But this is not a normal LSTM then, because it would mix also the other dimensions.\n\nI see some different definition in B.1. Is this just now a different notation, specifically in vector notation? Or is this really an alternative different model? As this says \"vector notation\", that really means that the initial equations (eq 1-7 and more) are all on scalars? Specifically, for example, eq 27 does not fit together with eq 4. If r is a vector in eq 4, then it just multiplies with a scalar h_{t-1} in eq 4, which is wrong.\n\nOr later in Sec 2.1, it then says \"In later formulations, multiple memory cells were combined in a vector, which allows the usage of recurrent weight matrices\". So actually only this aspect really leads to the vanilla LSTM. It means the presented formula in Sec 2.1 are really not the same as B.1. And also the Sec 2.2 formula are also not what you actually use, because then later you also have matrices R.\n\nI think this is all confusing. Why present the scalar variants at all, when they are never used like this?\n\n\nNow, for sLSTM, it was said that sLSTM has scalar memory and a scalar update. What does this mean? What is scalar about it? In what way is it different from the normal LSTM?\n\nI see in appendix B.2 that \"The matrices Rz , Ri , Rf , Ro are block-diagonal\". Is this specific only for the sLSTM? For the vanilla LSTM, you have fully dense matrices, right? So is this actually the main difference to the vanilla LSTM? Block-diagonal refers to \"scalar\"? So it's not really scalar?\n\n\n\n\nWhat is the motivation of exponential gates? To be able to open the gate and even amplify the input (gate > 1)? But why exp, why not relu or some other function which does not grow so fast?\n\nsLSTM Stabilized Version only in appendix. Why? Is this not so important?\n\nsLSTM normalization (n_t), how important?\n\n\n\n\nSec 2.2 \"we introduce a normalizer state that sums up the product of input gate times all future forget gates.\" - I think some word is missing.\n\n\nSec 2.3 \"we increase the LSTM memory cell from a scalar c ∈ R to a matrix C ∈ Rd×d\" - this is misleading, or I don't understand it. Actually you increase the LSTM memory cell from a vector to a matrix, not from a scalar to a matrix, or not? Or do you really have a different matrix C for every dimension, i.e. you actually have a 3D tensor as memory when there are multiple cells? Specifically, in eq 16-24, the i_t, f_t, h_t, q_t, etc, are those scalars or vectors now? In eq 16, when i_t/f_t are scalars, that means when you now have multiple cells, that the C_t becomes a 3D tensor?\n\n\nSec 2.4: xLSTM blocks and xLSTM architecture: I think it's bad to refer to the appendix for figures. I think those are crucial aspects about the whole model, which should be in the main text. Also, there should be some exact formulas to define the model, not just figures.\n\n\nSec 2.4 xLSTM architecture: \"LayerNorm ... See last column in Figure 5.\" - when I look at the last column in Fig 5, I only see some gray blocks. I don't see LayerNorm there.\n\n\nSec 2.4 xLSTM architecture: The only real \"definition\" of the model is Figure 5? In Figure 5, when it says xLSTM blocks, that means either a sLSTM block (like Figure 6) or a mLSTM block (like Figure 7)? So the light gray blocks are mLSTM, and the dark gray blocks are sLSTM in the last column in Fig 5?\n\nThe sLSTM and mLSTM blocks (Fig 6, Fig 7) have many further aspects which are not really discussed, like the group norm, etc. How did you end up with this specific architecture of the blocks? Have you tested any variations on that?\n\nSec 2.4: From Fig 5, it seems that the amount of sLSTM and mLSTM blocks is not the same? This is not really discussed or even mentioned at all? It seems to be also a crucial aspect of the whole xLSTM model. Only later in Sec 4, I read \"For all experiments, we use the notation xLSTM[a:b] for the ratio a/b of mLSTM-based versus sLSTM-based xLSTM blocks.\".\n\n\n\nFigure 4, scaling laws: I don't think it is reasonable to have the number of parameters on the x axis. More reasonable is to put e.g. the inference time at the x axis.\n\n\n\nRelation to multiplicative LSTM (https://arxiv.org/abs/1609.07959)?\n\nRelation to associative LSTM (https://arxiv.org/abs/1602.03032)?\n\nStudy on activation functions? Or only tanh?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces xLSTM, an advanced variant of traditional LSTMs, incorporating two extentions aimed at boosting its memory capacity and performance.\nThe first enhancement, termed sLSTM, modifies the standard LSTM by integrating exponential input and forget gates alongside a stabilizing normalizer term, designed to refine gate operations and ensure network stability.\n\nThe second mLSTM involves the incorporation of matrix-form hidden states into sLSTM, reminiscent of architectural advancements observed in models like the GLA, RetNet, and RWKV5/6. \nThis approach significantly augments the memory capacity of the LSTM units, aligning with recent trends in RNN design that prioritize enhanced memory retention.\n\nThe xLSTM architecture is constructed through a residual stacking of xLSTM token mixing layers coupled with Linear layers, mirroring the design principles employed in architectures such as Llama. \nThis configuration facilitates efficient information propagation and processing within the network.\n\nThe authors successfully scaled xLSTMs to 1B trained on hundreds of billions of tokens, from which they found xLSTMs perform better than exising Transformer, SSMs and other linear attention archs. \nAs for associative recall abilities, xLSTMs exibits better performance evaluated on MQAR tasks.\nThe paper also includes an exploration of scaling laws, revealing that xLSTMs maintain favorable performance characteristics when scaled to larger model sizes.\n\nSTRENGTHS:\nThe paper proposes a new RNN variant, denoted as xLSTM, making two extensions, i.e., sLSTM and mLSTM, to the classical LSTM, which has potentials to rejuvenate the classical RNNs in the LLM era.\nExperiements on 340M / 1B sized models show that xLSTMs perform very well on language modeling tasks.\nAnalysis indicates that the xLSTM variant excels in retrieval-intensive scenarios, a critical advancement given the acknowledged limitations in memory capacity that have historically constrained RNNs. I look forward to seeing the performence of xLSTMs scaled to larger sizes.\n\nWEAKNESSES:\nThe overall designs of mLSTM are not that novel. \nNotably, architectures such as the RetNet and GLA has successfully applied 2-d matrix-formed hidden states into RNN architectures that greatly enlarge the memory capabilities.\n\nxLSTMs additionally propose a exponentials gating, which, as claimed by the authors, can enable the abilities to solve state tracking problems. The authors verify the effectiveness of some synthetic tasks in Appendix D.1. \nHowever, I believe this design can greatly hurt the parallelizability of xLSTMs. \nTo enable stable training, the authors necessitate maintaining $i$ and $f$ in log-space first, at the scarifice of hardware optimizations, making it hard to benefit from tensor-core accelerations, a feature that is pivotal for achieving high computational efficiency. \nThis partially explains why their kernel impls are much slower than flash-attention.\nMoreover, the input and forget has to be kept in scalars, which I believe is inferior to the more fine-grained gating of GLA.\n\nQUESTIONS:\nIf possible, I would like to known the comparisons between xLSTMs and GLA\n\nLIMITATIONS:\nThe proposed exponential gating makes xLSTM hard to be parallelized, hindering its abilities to scale to larger sizes in compared to other exisiting architectures.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}