xLSTM: Extended Long Short-Term Memory

Summary

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

Strengths

I 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. Here are some of the key strengths of the paper: 1. An innovative approach to replacing input gates in LSTMs to improve storage decisions. 2. An intuitive methodology for storing and retrieving information in memory in matrix form. 3. Linear computation and constant memory complexity with respect to sequence length. 4. Extremely exhaustive evaluation in both language-based and synthetic tasks. [1] Gu, Albert, Karan Goel, and Christopher Ré. "Efficiently modeling long sequences with structured state spaces." arXiv preprint arXiv:2111.00396 (2021). [2] Orvieto, Antonio, et al. "Resurrecting recurrent neural networks for long sequences." International Conference on Machine Learning. PMLR, 2023.

Weaknesses

I think the paper has been executed almost flawlessly, so I will forgo this section and will simply add some questions in the section below.

Questions

For my interest, I would like the authors to clarify a few things: 1. 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? 2. 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? 3. 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? 4. 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? 5. 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. 6. 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)?

Rating

10

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Yes, the limitations of the proposed method are mentioned in the paper in a dedicated section.

Summary

Proposes some extension to LSTM. Specifically: * exponential gating with appropriate normalization and stabilization techniques * sLSTM with a scalar memory, a scalar update, and new memory mixing. This still needs to be calculated sequentially. * mLSTM that is fully parallelizable with a matrix memory (instead of just a vector) and a covariance update rule * putting sLSTM & mLSTM together into a residual block: the xLSTM block * Extended LSTM (xLSTM) blocks are residually stacked into xLSTM architecture Tested for language modeling. Compared to Transformers and State Space Models.

Strengths

Scales up LSTM models. Proposes extensions to the standard LSTM model which perform really well. Lots of experiments. (Although most is only in the appendix...) Open code.

Weaknesses

Some 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. There is never really an exact definition of the whole model. We can only infer it from the figures. (Or the code.) The 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.

Questions

In 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. However, 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. Or 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. I 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. Or 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. I think this is all confusing. Why present the scalar variants at all, when they are never used like this? Now, 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? I 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? What 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? sLSTM Stabilized Version only in appendix. Why? Is this not so important? sLSTM normalization (n_t), how important? Sec 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. Sec 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? Sec 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. Sec 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. Sec 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? The 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? Sec 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.". Figure 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. Relation to multiplicative LSTM (https://arxiv.org/abs/1609.07959)? Relation to associative LSTM (https://arxiv.org/abs/1602.03032)? Study on activation functions? Or only tanh?

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

4

Limitations

-

Summary

The paper introduces xLSTM, an advanced variant of traditional LSTMs, incorporating two extentions aimed at boosting its memory capacity and performance. The 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. The 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. This approach significantly augments the memory capacity of the LSTM units, aligning with recent trends in RNN design that prioritize enhanced memory retention. The 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. This configuration facilitates efficient information propagation and processing within the network. The 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. As for associative recall abilities, xLSTMs exibits better performance evaluated on MQAR tasks. The paper also includes an exploration of scaling laws, revealing that xLSTMs maintain favorable performance characteristics when scaled to larger model sizes.

Strengths

The 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. Experiements on 340M / 1B sized models show that xLSTMs perform very well on language modeling tasks. Analysis 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.

Weaknesses

The overall designs of mLSTM are not that novel. Notably, 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. xLSTMs 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. However, I believe this design can greatly hurt the parallelizability of xLSTMs. To 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. This partially explains why their kernel impls are much slower than flash-attention. Moreover, the input and forget has to be kept in scalars, which I believe is inferior to the more fine-grained gating of GLA.

Questions

If possible, I would like to known the comparisons between xLSTMs and GLA

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The proposed exponential gating makes xLSTM hard to be parallelized, hindering its abilities to scale to larger sizes in compared to other exisiting architectures.

please respond to rebuttal

Dear reviewers, The author-reviewer discussion period will end soon. Please make sure you read the authors' rebuttal and respond to it. If you have additional questions after reading the rebuttal please discuss with the authors. AC