PaperlayerView exact revision on arXiv

Source-linked AI summary

xLSTM: Extended Long Short-Term Memory

Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, Sepp Hochreiter

arXiv:2405.04517v2cs.LGcs.AIstat.ML

TL;DR

The paper asks how far LSTMs can scale for language modeling after addressing limitations in storage and revision. It introduces xLSTM through exponential gating, new scalar- and matrix-memory variants, and residual architectures, reporting favorable performance against Transformers and State Space Models. The authors conclude that xLSTM reaches at least the level of current technologies, while its matrix memory increases computational complexity.

  • Problem

    The paper investigates whether LSTMs can reach current large-language-model performance when their storage and storage-revision limitations are overcome and their scale is increased.

  • Method

    xLSTM combines exponential gating with sLSTM scalar memory and memory mixing, mLSTM matrix memory with covariance updates, and residual xLSTM blocks.

  • Results

    xLSTM performs favorably against Transformers and State Space Models in language modeling, with xLSTM[1:0] achieving lower perplexity in 568 of 571 PALOMA text domains.

  • Takeaways & Limitations

    Scaling laws indicate that larger xLSTM models could be serious competitors to current Transformer-based large language models.

  • Takeaways & Limitations

    mLSTM’s d×d matrix memory and updates increase computational complexity, although the computations can be parallelized on GPUs.

Abstract

from arXiv · show

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

LSTMs use a gated additive memory mechanism that supported broad sequence modeling success, but they face storage and revision limitations motivating xLSTM. The paper asks how far scaled LSTMs can go in language modeling after addressing these limitations.

  • LSTM foundations: LSTM updates its cell state additively, with input and forget gates moderating cell inputs and retained memory, while the output gate controls the hidden state.The cell state is normalized or squashed before producing the hidden state.
  • LSTM foundations: LSTMs supported text generation and other sequence-related tasks before Transformers became dominant in language modeling.Reported applications include handwriting generation, translation, program evaluation, and image captioning.
  • LSTM limitations: LSTMs struggle to revise a stored value when a later vector is more similar in nearest-neighbor search.The cited task scans a sequence for the most similar vector and returns its attached value at the end.
  • LSTM limitations: LSTMs have limited storage capacity because information must be compressed into scalar cell states, impairing rare-token prediction.The paper illustrates this limitation with perplexity on Wikitext-103 across token-frequency partitions.
  • Research question: The paper asks whether scaling LSTMs to current large-language-model sizes while overcoming these limitations can improve language modeling.This question follows the paper’s account of LSTM limitations and the emergence of Transformers.

2 Extended Long Short-Term Memory

xLSTM extends LSTM with exponential gating, new scalar and matrix memories, and residual architectures designed to improve storage, parallelization, and scalability.

  • Extended LSTM: xLSTM introduces exponential gating with normalization and stabilization, plus sLSTM and mLSTM memory structures.sLSTM uses scalar memory, scalar updates, and memory mixing; mLSTM uses matrix memory and covariance updates.
  • sLSTM: Exponential gating enables LSTM storage decisions to be revised when more similar information appears later in a sequence.A normalizer state tracks gated contributions, supporting stable retrieval.
  • sLSTM: sLSTM supports multiple memory cells and heads with recurrent memory mixing within heads and across cells.This mixing distinguishes sLSTM from mLSTM, whose recurrence omits hidden-hidden memory mixing for parallelization.
  • mLSTM: mLSTM expands scalar memory c ∈ R into matrix memory C ∈ R^d×d for key-value storage and covariance-based retrieval.Its gates correspond to decay, learning rate, and output scaling, while the recurrence can be reformulated in parallel.
  • xLSTM Architecture: xLSTM blocks integrate sLSTM or mLSTM into residual post- or pre-up-projection designs, then architectures residually stack these blocks.The sLSTM block uses post up-projection; the mLSTM block uses pre up-projection.
  • Memory and Speed: xLSTM has linear computation and constant sequence-length memory complexity, while mLSTM trades higher matrix-memory computation for greater capacity and GPU parallelism.mLSTM’s d×d memory and update are computationally expensive, although parallel execution can limit wall-clock impact.

3 Related Work

Related work addresses Transformer complexity and recurrent sequence modeling through linear attention, state space models, gated RNNs, covariance updates, and residual architectures.

  • Linear Attention: Linear-attention methods approximate or reformulate attention to reduce its quadratic context-length complexity.Examples include Synthesizer, Linformer, Linear Transformer, and Performer.
  • State Space Models: State Space Models provide linear context-length computation and have shown promising performance compared with Transformers.The cited progression includes S4, DSS, GSS, S5, BiGS, and H3.
  • Recurrent Neural Networks: Recent recurrent models such as LRUs, HGRN, HGRN2, and RWKV target Transformer replacement while retaining linear context-length computation.These approaches include recurrent architectures developed for language modeling.
  • Gating: Gating appears across HGRN, GLA, GSS, BiGS, MEGA, RWKV, and Mamba as a recurring design idea in modern sequence models.The passage frames gating as an LSTM idea that was rediscovered and reinterpreted.
  • Covariance Update Rule: Covariance-style updates for memory appear in Fast Weight Programmers, RWKV-5/6, Retention, Linear Transformer, and HGRN2.xLSTM incorporates this mechanism into the mLSTM matrix-memory framework.
  • Most Related: xLSTM is conceptually closest to Retention, RWKV, and HGRN2, but differs from them through sLSTM memory mixing.The paper connects memory mixing with state-tracking capability and greater expressivity than SSMs and Transformers.
  • Residually Stacking Architectures: Like contemporary large deep learning models, xLSTM architectures use residual stacking, a construction also used by deep convolutional networks and Transformers.The architecture relies on pre-LayerNorm residual backbones.

4 Experiments

The experiments evaluate xLSTM on synthetic memory and state-tracking tasks, long-context benchmarks, language modeling, downstream reasoning, scaling, and generation efficiency. Across these evaluations, xLSTM generally performs strongly, including favorable results against Transformers and State Space Models.

  • Experimental scope: xLSTM is evaluated on formal languages, associative recall, Long Range Arena, 15B-token language modeling, ablations, and 300B-token large-model experiments.The experiments cover both targeted capabilities and broad language-modeling comparisons.
  • Synthetic Tasks and Long Range Arena: xLSTM performs best among non-Transformer models on Multi-Query Associative Recall, including the most difficult setting with 256 key-value pairs.The task uses context length 2048 and tests increasing numbers of key-value pairs.
  • Synthetic Tasks and Long Range Arena: xLSTM demonstrates consistent strong performance across Long Range Arena tasks, suggesting efficiency in handling different long-context problems.The evaluation compares xLSTM with existing methods on long-sequence tasks.
  • Method Comparison and Ablation Study: On 15B SlimPajama tokens, xLSTMs with new memory achieve the best validation perplexities, while ablations attribute improvement over vanilla LSTM to exponential gating and matrix memory.The gating ablation also finds incremental gains when gates are learnable and input-dependent.
  • xLSTM as Large Language Model: At 1.3B scale, xLSTM maintains low perplexities when extrapolating from context length 2048 to 16384, unlike the other compared methods.The comparison includes RWKV-4, Llama, and Mamba.
  • xLSTM as Large Language Model: Across model sizes, xLSTM is best on validation perplexity and the vast majority of downstream tasks, with Mamba sometimes best only on ARC.The evaluation uses models trained on 300B SlimPajama tokens.
  • xLSTM as Large Language Model: On PALOMA, xLSTM[1:0] has lower perplexity than Mamba on 568 of 571 domains, Llama on 486, and RWKV-4 on 570.These correspond to 99.5%, 85.1%, and 99.8% of domains, respectively.
  • xLSTM as Large Language Model: Scaling-law results show similar scaling behavior with different offsets, with xLSTM ahead of Mamba and projected to remain favorable at larger sizes.RWKV-4 performs worst, followed by Llama and Mamba.

5 Limitations

The paper identifies computational, memory, initialization, and scaling limitations affecting the current xLSTM implementations and longer contexts.

  • sLSTM memory mixing is not parallelizable, although a CUDA kernel makes it less than two times slower than parallel mLSTM.The mLSTM implementation is currently about four times slower than FlashAttention or the scan used in Mamba.
  • mLSTM CUDA kernels are not optimized, and its matrix memory requires processing d × d matrices.Despite the computation complexity, the memory update and retrieval can use parameter-free standard matrix operations, keeping wall-clock overhead minor.
  • Forget-gate initialization must be chosen carefully, while matrix memory may overload longer contexts because it is independent of sequence length.The paper reports no apparent limitation for contexts up to 16k.
  • Large language models face expensive computational loads in the current xLSTM implementation.

6 Conclusion

The paper extends LSTMs with exponential gating and new memory structures, integrates them into residual architectures, and reports favorable language-modeling performance relative to Transformers and State Space Models.

  • 6 Conclusion: xLSTM models perform favorably in language modeling compared with state-of-the-art Transformers and State Space Models.
  • 6 Conclusion: The scaling laws indicate that larger xLSTM models could become serious competitors to Transformer-based Large Language Models.
  • 6 Conclusion: The parallel mLSTM forward pass processes all timesteps of a full sequence simultaneously, while recurrent formulation remains available for fast text generation.

B Experiments

The experiments used modern GPU software and substantial cluster compute, including a nearest-neighbor task and 125M-parameter rare-token prediction models.

  • Experiments used Python 3.11, PyTorch 2.2.0, CUDA 12.1, and NVIDIA A100 GPUs across a 128-node cluster.Each node contained eight NVIDIA A100 GPUs, and training occurred over three months.
  • The nearest-neighbor experiment trained two-block models with embedding dimension 128 on sequences up to length 64.Reported values were means for the best learning rate with 99% confidence intervals.
  • Rare-token prediction used 125M-sized models trained on Wikitext-103, with Llama, LSTM, and mLSTM-based xLSTM configurations.

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

The formal-language evaluation probes generalization across Chomsky-hierarchy tasks and extrapolates beyond training lengths; xLSTM performs best among evaluated models where no model achieves proper extrapolation.

  • B.1.1 Test of xLSTM’s Exponential Gating with Memory Mixing.: Formal-language tasks test xLSTM’s exponential gating and memory-mixing mechanism across different expressivity levels.The tasks are grouped by the Chomsky hierarchy and include counting and memory-focused problems.
  • B.1.1 Test of xLSTM’s Exponential Gating with Memory Mixing.: Models train on task lengths up to 40 and evaluate on lengths from 40 to 256 to measure task generalization.
  • B.1.1 Test of xLSTM’s Exponential Gating with Memory Mixing.: The comparison includes Llama, Mamba, Retention, Hyena, RWKV variants, LSTM, and three xLSTM configurations.
  • B.1.1 Test of xLSTM’s Exponential Gating with Memory Mixing.: Although no model achieves proper extrapolation to larger context lengths, xLSTM performs best among the evaluated models.This result concerns supplementary tasks where no model reached a minimum scaled accuracy of 0.3.
  • B.1.1 Test of xLSTM’s Exponential Gating with Memory Mixing.: The task suite includes missing-duplicate, modular-arithmetic, odds-first, parity, repetition, set, and equation-solving problems.

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

The experiments test xLSTM memory on Multi-Query Associative Recall, varying context length and key-value pairs across several recurrent and sequence models. xLSTM[1:1] performs best in harder settings and maintains performance during context-length extrapolation.

  • Experiment setup: MQAR tests associative recall by requiring models to retrieve values associated with keys across varying context lengths and key-value counts.The experiments use two blocks for each model and vary context length, key-value pairs, and model size.
  • Experiment 1: xLSTM[1:1], xLSTM[1:0], Mamba, RWKV-5, and RWKV-6 solve the smaller Experiment 1 settings across all model dimensions, unlike H3, Hyena, and RWKV-4.The comparison includes Llama, Retention, H3, Hyena, several RWKV variants, LSTM variants, and xLSTM variants.
  • Experiment 1: Exponential gating improves associative-recall capability, while xLSTM[1:1] and xLSTM[1:0] provide greater memory capacity than xLSTM[0:1] and original LSTM variants.The comparison of xLSTM[0:1] with both original LSTM variants indicates a gating benefit, but those models remain limited by reduced memory capacity.
  • Experiment 2: Increasing key-value pairs affects performance more than increasing context length, and xLSTM[1:1] is the best-performing non-Transformer model in Experiment 2.RWKV variants lose performance with longer contexts, while xLSTM[1:1] and xLSTM[1:0] show no clear decay.
  • Experiment 3: In Experiment 3, xLSTM[1:1], xLSTM[1:0], and Mamba show unchanged performance during extrapolation, while RWKV models degrade slightly as context length increases.xLSTM[1:1] remains best because it maintains its superior Experiment 2 performance.

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

The Long Range Arena evaluation examines xLSTM across five tasks involving long contexts and diverse data types. xLSTM demonstrates consistent strong performance across all evaluated tasks.

  • Evaluation: The Long Range Arena evaluation tests retrieval, ListOps, image, and other tasks spanning longer contexts and diverse data types.The study omits text classification and Pathfinder-X because language modeling and the original Pathfinder cover those settings elsewhere.
  • Evaluation: The comparison includes Llama, Mamba, LSTM, RWKV-4, and xLSTM, selecting the better xLSTM[0:1] or xLSTM[1:0] variant by validation accuracy.The selected xLSTM variant is xLSTM[0:1] for image tasks and xLSTM[1:0] for all other tasks.
  • Results: xLSTM demonstrates consistent strong performance on all Long Range Arena tasks.Table 6 marks the best model in bold and the second-best underlined; xLSTM denotes the validation-selected variant.

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

The SlimPajama study compares xLSTM with other sequence architectures and ablates its components under a shared language-model training setup. The ablations identify residual structure, exponential gating, matrix-memory placement, convolution, and gate design as important choices.

  • Method comparison: The SlimPajama comparison trains models with context length 2048 and evaluates perplexity under a common 15B-token setup.The study uses AdamW, shared tokenizer and training conventions, and reports validation perplexity for ablations.
  • Method comparison: A repository learning rate gives RWKV-6 perplexity 15.03, outperforming its 4e-4 learning-rate result of 16.38.The paper reports that repository learning rates also performed worse than the rates used for xLSTM across other model sizes.
  • Matrix-memory ablation: Matrix memory performs best in the pre up-projection block, and the dimension-wise causal convolution in that block is important.The comparison uses xLSTM[1:0] and tests matrix-memory placement with and without the convolution.
  • Component ablation: Adding skip-connections and pre-LayerNorm enables training large vanilla LSTMs, while a post up-projection backbone and exponential gating further improve performance.Replacing every second LSTM layer with a GeLU feed-forward network boosts performance before exponential gating yields the sLSTM.
  • Gating ablation: Learnable, input-dependent gates produce the best mLSTM results among the gating configurations tested.The paper contrasts this design with other methods using related gating approaches.
  • Scaling: Scaling experiments include xLSTM models from 125M through 2.7B parameters, trained with context length 2048.The supplied setup specifies batch size 256 and peak learning rates from Table 7 for the scaling experiments.

C Detailed Results on PALOMA Language Model Evaluation

The PALOMA evaluation reports detailed perplexities for 1.3B-parameter models trained on 300B SlimPajama tokens across 571 subdomains.

  • PALOMA evaluation: 1.3B models trained on 300B SlimPajama tokens are evaluated on 571 PALOMA subdomains with detailed perplexity values.The aggregated perplexities in Table 4 are not macro averages of the subdomain values in Table 10.
Loading 2405.04517v2…