[Paper Review] xLSTM: Extended Long Short-Term Memory

Maximilian Beck, Korbinian Pöppel+7 more|arXiv (Cornell University)|May 7, 2024

Topic Modeling86 citations

TL;DR

xLSTM introduces exponential gating and two memory-augmented variants (sLSTM with scalar memory and mLSTM with matrix memory) to overcome LSTM limitations, achieving favorable language modeling performance and scaling compared to Transformers and State Space Models.

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.

Motivation & Objective

• Motivate the need to overcome classic LSTM limitations (revision of storage, limited memory, lack of parallelism) in large-scale language modeling.
• Propose two memory-augmented LSTM variants (sLSTM and mLSTM) with exponential gating.
• Introduce xLSTM blocks that integrate these variants into residual architectures for scalable modeling.
• Demonstrate through synthetic tasks and large-scale language modeling experiments that xLSTM can rival or surpass Transformer/SSM baselines in certain settings.

Proposed method

• Introduce exponential gating with normalization and stabilization to enable revisable storage decisions.
• Develop sLSTM with a scalar memory, scalar update, and memory mixing across cells/heads.
• Develop mLSTM with a matrix memory and a covariance update rule to enable high-capacity storage and parallelizable retrieval.
• Create xLSTM blocks by embedding sLSTM or mLSTM inside residual block backbones (post or pre up-projection depending on variant).
• Stack xLSTM blocks into xLSTM architectures using pre-LayerNorm backbones for scalable modeling.
• Provide a CUDA-optimized implementation and discuss memory/compute characteristics relative to Transformers.

Experimental results

Research questions

• RQ1 Can exponential gating and new memory structures remediate fundamental LSTM limitations (revision of storage, memory capacity, and sequential bottlenecks) at scale?
• RQ2 How do sLSTM and mLSTM compare to Transformers and State Space Models in language modeling tasks when scaled to billions of parameters?
• RQ3 Do xLSTM architectures with residual stacking deliver competitive perplexities and downstream task performance across model sizes and data scales?

Key findings

• Exponential gating and matrix/scalar memory designs enable xLSTM to outperform several baseline models on language modeling tasks in validation perplexity.
• In synthetic and long-range tasks, xLSTM demonstrates improved state-tracking and memory capacity over traditional LSTM and some Transformer/SSM variants.
• xLSTM[1:0] and xLSTM[7:1] achieve strong validation perplexities across model sizes in SlimPajama experiments and show favorable scaling behavior.
• On 300B-token training, xLSTM variants maintain better sequence-length extrapolation, longer-context perplexities, and competitive downstream task performance compared to RWKV, Llama, and Mamba.
• Ablation studies indicate that exponential gating and matrix memory are primary drivers of performance gains, with learnable, input-dependent gates providing additional benefits.
• xLSTM architectures are capable of competitive performance with linear-like memory considerations and parallelizable components, offering a scalable alternative to Transformer-dominated regimes.