xLSTM: Extended Long Short-Term Memory

In plain words

LSTMs are recurrent networks that have been used for sequences such as text. Transformers are the main comparison in this paper. The paper asks how far LSTMs go when they are scaled up with techniques from modern language models and when their known weak points are fixed.

LSTMs have three weak points. They struggle to replace a stored value when a better match shows up later in the sequence. Everything they remember has to fit into a small number of scalar slots. And their step-by-step recurrence is hard to spread across a sequence during training. The authors change two things. First, the gates can be exponential, so the model can strongly overwrite or boost what it stores. Extra bookkeeping keeps the numbers from overflowing. Second, the memory can be a small matrix that works like a table of key-value pairs. A value is written by adding the outer product of a key and a value, and read back with a query.

These two changes give two new cells. The sLSTM keeps a scalar memory and lets units within the same head influence each other. The mLSTM has a matrix memory and no link from the previous output into the gates, so it can be computed in parallel over the sequence. Both cells sit inside residual blocks, and stacking those blocks gives the xLSTM model. The authors report that the resulting language models compete with, or beat, Transformers and state-space models.

TL;DR

• xLSTM extends the LSTM with exponential input and forget gates, a normalizer state and a stabilizer state §2.
• sLSTM is a scalar-memory cell with memory mixing inside each head; mLSTM is a matrix-memory cell with a covariance update and no hidden-to-hidden recurrence, so it can be parallelised §2.2, §2.3.
• On SlimPajama with 15B tokens, xLSTM[1:0] reaches 13.43 validation perplexity against 13.70 for Mamba and 14.25 for Llama tab. 1. At 300B tokens, 1.3B-parameter xLSTM[1:0] reaches 8.89 against 9.14 for Mamba and 9.44 for Llama tab. 3. All numbers are author-reported.
• The released code is Apache-2.0. Its gate inputs, output gate and normaliser differ from the paper's equations in several places (see "Paper vs code").

Problem

LSTMs have three limitations §1.

• Inability to revise storage decisions. In a nearest-neighbour search over a sequence, an LSTM struggles to replace a stored value when a more similar vector appears later fig. 2, left.
• Limited storage capacity. Information must be compressed into scalar cell states. On WikiText-103, LSTM perplexity is worse on rare tokens than on frequent ones fig. 2, right.
• No parallelisation. Hidden-to-hidden connections (memory mixing) force sequential processing §1, §2.5.

The paper asks what a scaled-up LSTM with these limitations removed can reach in language modelling, compared with Transformers and state-space models §1.

Method

The model keeps the LSTM's constant error carousel: an additive cell-state update that is gated §2.1. Four changes follow.

1. Exponential gating with normalizer and stabiliser. The input gate, and optionally the forget gate, use the exponential function. A normalizer state sums the input-gate-weighted keys or cell inputs, scaled by the future forget gates, and keeps track of the gates' overall strength. A stabilizer state m_t is subtracted inside the exponentials so their arguments stay at or below zero. The stabilised and unstabilised forms give identical outputs and gradients with respect to the parameters §2.2; app. A.2. The paper allows either a sigmoid or an exponential forget gate §2.2, §2.3.

2. sLSTM (scalar memory, memory mixing) §2.2; app. A.2. The cell update is the LSTM update with exponential gates and a normaliser. Memory mixing uses recurrent matrices from the previous hidden state to all gate pre-activations. These matrices are block-diagonal, one block per head, so there is no mixing across heads §2.2; app. A.2. The cell input activation is tanh. The hidden state is the output gate times the normalised cell app. A.2.

3. mLSTM (matrix memory, covariance update) §2.3; app. A.3. The cell is a d x d matrix C. Keys and values are written with the covariance rule, and gate-weighted retrieval uses the query. A normaliser vector n accumulates keys. The retrieved vector is divided by max(|n·q|, 1) in the main text, or by max(|n·q|, exp(-m)) in the stabilised form of the appendix. An output gate scales the result. The input and forget gates are linear in the input x_t. The keys carry a 1/sqrt(d) factor. A parallel form over the whole sequence is given for training app. A.3.

4. Blocks and architecture §2.4; fig. 3; app. A.4. An sLSTM block is a residual block with a post up-projection: the sLSTM output is followed by a gated GeLU MLP. An mLSTM block is a residual block with a pre up-projection: the input is first projected up by a factor of 2 and the mLSTM sits inside, with a causal convolution, a learnable skip connection, a group-normalised output and an output gate. Blocks use pre-LayerNorm and are stacked residually. The notation xLSTM[a:b] gives the ratio of mLSTM to sLSTM blocks §2.4; §4.

Key equations, in the paper's notation (reproduced in our own symbols):

C_t = f_t\, C_{t-1} + i_t\, v_t k_t^\top, \qquad n_t = f_t\, n_{t-1} + i_t\, k_t \qquad [§2.3]

\tilde h_t = \frac{C_t q_t}{\max\{|n_t^\top q_t|, 1\}}, \qquad h_t = o_t \odot \tilde h_t \qquad [§2.3]

m_t = \max(\log f_t + m_{t-1},\ \log i_t) \qquad [§2.2, \text{stabilizer}]

Pseudocode for one mLSTM step (one head), following the stabilised form in the appendix:

# One mLSTM step for one head  [§2.3] [app. A.3]
# State: C (d, d) matrix memory, n (d,) normalizer, m scalar stabilizer
def mlstm_step(x, C_prev, n_prev, m_prev):
    # Query, key and value projections; keys are scaled by 1/sqrt(d)  [§2.3]
    queries = W_q @ x + b_q                       # (d,)
    keys = (W_k @ x + b_k) / sqrt(d)              # (d,)
    values = W_v @ x + b_v                        # (d,)
    # Gate pre-activations from the input  [§2.3]
    i_tilde = w_i @ x + b_i                       # scalar
    f_tilde = w_f @ x + b_f                       # scalar
    o_tilde = W_o @ x + b_o                       # (d,)
    # Stabilizer: keeps every exp() argument at or below zero  [app. A.2]
    log_f = log_sigmoid(f_tilde)
    m = max(log_f + m_prev, i_tilde)
    i_gate = exp(i_tilde - m)
    f_gate = exp(log_f + m_prev - m)
    # Covariance update of the matrix memory and of the normalizer  [§2.3]
    C = f_gate * C_prev + i_gate * outer(values, keys)
    n = f_gate * n_prev + i_gate * keys
    # Retrieval, divided by max(|n . q|, exp(-m)) plus a small eps  [app. A.3]
    normalizer = max(abs(dot(n, queries)), exp(-m))
    h_tilde = (C @ queries) / (normalizer + eps)
    # The output gate scales the retrieved vector  [§2.3]
    h = sigmoid(o_tilde) * h_tilde
    return h, C, n, m

Pseudocode for one sLSTM step (one head, sigmoid forget gate shown; the paper also allows an exponential one):

# One sLSTM step for one head  [§2.2] [app. A.2]
# State: h, c, n (each (d,)) and m (d,) stabilizer
# R_* are per-head recurrent matrices (memory mixing)  [app. A.2]
def slstm_step(x, h_prev, c_prev, n_prev, m_prev):
    # Pre-activations: input, memory mixing and bias terms  [§2.2]
    z_tilde = W_z @ x + R_z @ h_prev + b_z        # (d,)
    i_tilde = W_i @ x + R_i @ h_prev + b_i        # (d,)
    f_tilde = W_f @ x + R_f @ h_prev + b_f        # (d,)
    o_tilde = W_o @ x + R_o @ h_prev + b_o        # (d,)
    # Cell input activation is tanh; output gate is sigmoid  [app. A.2]
    z = tanh(z_tilde)
    o_gate = sigmoid(o_tilde)
    # Stabilizer, then exponential input gate and forget gate  [§2.2]
    log_f = log_sigmoid(f_tilde)
    m = max(log_f + m_prev, i_tilde)
    i_gate = exp(i_tilde - m)
    f_gate = exp(log_f + m_prev - m)
    # Cell and normalizer updates  [§2.2]
    c = f_gate * c_prev + i_gate * z
    n = f_gate * n_prev + i_gate
    # Hidden state: output gate times the normalised cell  [§2.2]
    h = o_gate * (c / n)
    return h, c, n, m

Pseudocode for one residual block of the stack:

# One xLSTM block: pre-LayerNorm residual  [§2.4] [fig. 3] [app. A.4]
# layer is an mLSTM layer or an sLSTM layer; ffn is a gated MLP
def xlstm_block(x, layer, ffn, is_slstm_block):
    # Residual sequence mixer  [§2.4]
    x = x + layer(layer_norm(x))
    # Only sLSTM blocks add a gated MLP after the mixer  [§2.4] [app. A.4]
    if is_slstm_block:
        x = x + ffn(layer_norm(x))
    return x

Hyperparameters

Values for the released blocks, the ablations and the large-scale runs. "not stated" means the paper does not give the value.

Results

All numbers are author-reported. Language-model numbers are validation perplexities on SlimPajama (lower is better), with the paper's setup in the caption of each table.

15B tokens, models matched to a 350M GPT-3 size (dim 1024, 24 blocks) §4.2; tab. 1

Ablation on the 15B setting §4.2; tab. 2

Top part (vanilla LSTM to xLSTM):

Bottom part (gate variants for xLSTM[1:0]; the table gives the forget gate and input gate settings: input-dependent, learnable bias, bias init):

Block variants for the matrix memory app. B.2, table "Matrix Memory variants": post up-projection with ReLU^2 and no gating, 24 blocks, 430.4M: 13.90. Pre up-projection without the convolution, 48 blocks, 408.8M: 15.41. Pre up-projection with the convolution, 48 blocks, 409.3M: 13.43 (the value the paper uses for xLSTM[1:0]).

300B tokens, 125M to 1.3B §4.3; tab. 3

Downstream tasks are scored with the LM Evaluation Harness: accuracy for LAMBADA, WinoGrande, PIQA and ARC-E; length-normalised accuracy for HellaSwag and ARC-C app. B.3. Each model reports its best of fp32, bf16 and mixed precision app. B.3.

PALOMA, average perplexity over the grouped domains §4.3; tab. 4 (the table has 12 further per-domain columns, not reproduced here):

Across 571 PALOMA domains, xLSTM[1:0] has lower perplexity than Mamba in 568, than Llama in 486 and than RWKV-4 in 570 §4.3.

Other results

• Formal languages §4.1; fig. 4; app. B.1.1. Two-block models are trained on task lengths up to 40 and evaluated on lengths 40 to 256. Transformers and SSMs without memory mixing cannot solve parity. xLSTM does best among the tested models on tasks that no model extrapolates well.
• Multi-query associative recall §4.1; fig. 5; app. B.1.2. Key-value pairs go up to 256, context up to 2048. xLSTM[1:1] is the best non-Transformer model, including at small sizes fig. 5. The sLSTM block does not reduce memory capacity §4.1.
• Long Range Arena (test accuracy) app. B.1.3; tab. in app. B.1.3:

"X" marks a model that does not beat the random baseline. For each xLSTM task the better of xLSTM[1:0] and xLSTM[0:1] was picked on validation accuracy app. B.1.3.

• Sequence extrapolation §4.3; fig. 7. 1.3B models trained at context 2048 and tested up to 16384 tokens. Perplexity at 16k: xLSTM[7:1] 8.92, xLSTM[1:0] 9.01, RWKV-4 13.75, Mamba 14.00, Llama 337.83.
• Scaling laws §4.3; fig. 6, fig. 8. Over 125M to 1.3B, all models follow similar power laws with different offsets. RWKV-4 is worst, then Llama and Mamba. xLSTM is better than Mamba. The authors extrapolate this to larger sizes. The 15B-token scaling plot shows the xLSTM[7:1] point at 2.7B degraded by slower early convergence fig. 6 caption.
• Generation speed §2.5, §4.3; fig. 9. At 1.3B with a 16-token prefill and batch size 1, generation time grows linearly for xLSTM, Mamba and RWKV-4, and quadratically for Llama. For decoding throughput on an A100-80GB, Llama runs out of memory already at small batch sizes, while xLSTM[1:0] sustains higher throughput than Mamba fig. 9.

Paper vs code

Code: the official repository at commit ab22eadbd245f293dd8dec38ed29963d73758a12 (Apache-2.0). It is linked from the abstract. The repository contains only the parity-task experiment configs (experiments/parity_*.yaml). The SlimPajama training configs are not included.

• mLSTM gate inputs. The paper's gates are linear in x_t §2.3. The code computes all four mLSTM gate pre-activations from the concatenation of q, k and v (3 x inner dimension) cell.py, with the input gate and forget gate weights defined at cell.py.
• mLSTM output gate and skip. The paper's output gate is a sigmoid of a linear map of x_t §2.3. The code's output gate is SiLU applied to one half of the up-projection layer.py. The code also adds a learnable skip connection from the convolved q/k activations before the gate, and applies a head-wise group norm to the cell output layer.py. The paper mentions the skip and the group norm only in the block figure app. A.4.
• Normaliser threshold. The main text uses max(|n·q|, 1) §2.3. The code uses max(|n·q|, exp(-m_t)) plus an eps of 1e-6 added to the denominator backends.py. The two are equal only when m_t = 0.
• Forget gate. The paper allows a sigmoid or an exponential forget gate §2.2, §2.3, app. A.1. Both the parallel and the recurrent mLSTM backends apply log-sigmoid to the forget gate and nothing else backends.py. The sLSTM vanilla pointwise function does the same slstm.py. The CUDA kernels were not read for this card.
• Parallel form. The paper describes the parallel form as the training form app. A.3. The code's default mLSTM backend is the parallel one cell.py. It builds per-head T x T matrices backends.py. The code also has a chunkwise variant backends.py that the paper does not describe. The cell's default is the parallel one.
• sLSTM gate clamp and first-step special case. The vanilla pointwise code caps the stabilised input and forget gates at 1 slstm.py. Because the stabiliser is a maximum, the cap never binds. The code also special-cases a zero normaliser by setting the stabiliser to the input pre-activation slstm.py. The paper does not mention either.
• sLSTM bias initialisation. The paper gives the [3,6] forget-bias initialisation only for mLSTM tab. 2. The sLSTM cell defaults to a block-dependent power-law forget-bias initialisation and zero recurrent weights cell.py. The linear [3,6] option exists as small_init cell.py.
• Recurrent-gradient clipping. The appendix sets the clip to 10 app. A.2. The sLSTM cell's default is None cell.py, and no repository config sets it.
• sLSTM backend. The cell defaults to the CUDA backend cell.py with bf16 recurrent weights and fp32 biases cell.py. The vanilla Python path is the reference implementation.
• sLSTM convolution and gate names. The convolution feeds only the input and forget gates, and z and o use the raw input layer.py. This matches the appendix figure caption app. A.4. The layer names are crossed: the module called fgate produces the input-gate pre-activation and igate the forget-gate one layer.py. The model is consistent, because the outputs are concatenated in the order the cell expects, but the names mislead.
• Feed-forward size. The sLSTM block MLP is 4/3 in the figure caption app. A.4. The code's default is 1.3 feedforward.py. The paper's SlimPajama configs are not in the repository, so the size used for the LM runs cannot be checked.
• Block structure. An sLSTM block is a pre-LayerNorm residual sLSTM mixer followed by a gated GeLU MLP xlstm_block.py. An mLSTM block is the mixer alone, with its own up-projection and no separate MLP block.py. This matches §2.4.
• Stack and head defaults. The post-block LayerNorm is on by default xlstm_block_stack.py. The LM head has no bias and is not tied by default xlstm_lm_model.py. The embedding is not weight-decayed by default xlstm_lm_model.py. The paper does not state these defaults.

Pitfalls

• Training cost is quadratic in sequence length. The paper says xLSTM has linear computation and constant memory in sequence length §2.5. That holds for recurrent decoding backends.py. The parallel training form used by default builds T x T matrices per head backends.py. Budget training memory accordingly. The paper says the mLSTM CUDA kernel is about four times slower than FlashAttention or Mamba's scan §5 (ii).
• The stabiliser must be implemented exactly. The exponentials in both gates overflow without it §2.2; app. A.2. Use the subtraction of the running maximum, and keep the normaliser threshold at exp(-m_t) rather than 1 backends.py.
• Forget-gate initialisation is critical. The paper lists this as a limitation §5 (iv). The mLSTM forget bias is spread linearly from 3 to 6, which starts the forget gate near 1, and the input bias is near zero cell.py. The sLSTM default differs (see "Paper vs code").
• Gates from q, k and v. An implementation that follows the paper's gate equations, with gates linear in x_t, will not match the released code. Pick one and say which.
• Output-gate activation. A sigmoid output gate (paper) and a SiLU gate with group norm and skip (code) are different models. Decide which you follow.
• Forget gate choice. The paper allows an exponential forget gate. The released backends use only the sigmoid. The paper's stabiliser equations are written in terms of log f, so they cover both variants, but the code never runs the exponential one.
• Recurrent-gradient clip. The paper clips the recurrent gradient at 10 app. A.2. Without that clip, sLSTM can give exploding gradients on long contexts app. A.2. The repo does not set it. Test with and without.
• Reported numbers are re-evaluations for some rows. The source tables mark the 13.43 and 13.48 values as re-evaluations. The comments in the same files list training-time values for the same runs that differ slightly (for example 13.50 for an xLSTM[7:1]-type run).
• Table head dimensions use the inner size. The "#Heads / Head dim" column in the paper's scaling table gives the mLSTM head dimension after the factor-2 up-projection app. B.2; code: xlstm/blocks/mlstm/layer.py#L23.
• The 15B comparison uses one LR. All models use peak LR 1e-3 in the 15B comparison, not the per-model values of the scaling table app. B.2.
• Missing LM configs. The SlimPajama LM training configs are not in the repository. Hyperparameters for the LM runs come from app. B.2 and app. B.3 only. Several details are not stated: data version and preprocessing, sLSTM gate initialisation for the LM runs, and the exact FFN size.
• Limits the paper states. The sLSTM is not parallelisable, but its fast CUDA kernel is less than two times slower than the parallel mLSTM §5 (i). The d x d matrix memory has high computation cost §5 (iii). The matrix memory does not depend on sequence length, so long contexts can overload memory §5 (v). The paper did not tune the large architectures or their hyperparameters §5 (vi).
• Compute. Experiments used 128 nodes of eight A100 GPUs for about three months, with more than 95% spent on language modelling app. B.

Lineage

• Builds on: the original LSTM with constant error carousel and gating §1, §2.1; the forget gate introduced by Gers et al. §2.1; the memory-mixing LSTM variant of Greff et al. §2.1; app. A.1; the covariance update rule, as cited in §2.3; Fast Weight Programmers, cited as Schmidhuber et al. and Schlag et al. §2.3; the linear Transformer of Katharopoulos et al. §3. The paper cites these by name only, without arXiv ids.
• Compared against: GPT-3 and Llama (Transformers); H3, Mamba (SSMs); RWKV-4, RWKV-5, RWKV-6, HGRN, HGRN2 (RNNs); GLA, RetNet, Hyena (linear Transformers) §4.2; Llama, Mamba and RWKV-4 at 300B tokens §4.3.
• Superseded by: not stated in this paper.