xLSTM for fast inference

mLSTM cell
$x_t$	input vector at token $t$
$q_t,$ $k_t,$ $v_t$	query, key and value projections of $x_t$
$\phi$	feature map applied to queries and keys
$S_t,$ $z_t$	matrix memory and normalizer, fixed size
$i_t,$ $f_t,$ $o_t$	input, forget and output gates
$h_t$	cell output at token $t$
Distillation
$\hat h_t$	output of the hybrid attention layer
$p_T,$ $p_\theta;$ $\theta$	teacher and student next-token distributions; student parameters
$y_t;$ $x_{1:t}$	next token; the text up to $t$
$\gamma,$ $\beta$	weights of the cross-entropy and KL terms (0.9 and 0.1)
$\mathcal B;$ $A_S(b),$ $A_T(b)$	benchmark set; student and teacher scores on benchmark $b$
$\alpha;$ $C_\alpha;$ $\alpha^*$	tolerance; win-and-tie rate; smallest tolerance with $C_\alpha\ge1/2$

Context. A transformer decodes by attending to a KV cachethe stored keys and values of every earlier token, read in full at each decoding step; it grows by one key and one value per token, per head and per layer that grows with every token, so the time and memory of each step grow with the context. That matters for a robot controller, which must answer at the rate of its sensors: 100 to 1,000 Hz, or under 10 ms per step, and 15 seconds of history at 1,000 Hz is a context of 15,000 steps. It matters in serving language models too, where the cache limits batch size and long generations. xLSTMa recurrent architecture (Beck et al., 2024) built from mLSTM blocks, a gated linear attention with a matrix memory, and sLSTM blocks, a scalar-memory LSTM with recurrent mixing and Mambaa state-space sequence layer (Gu and Dao, 2023) with input-dependent dynamics, trained in parallel and run as a recurrence at inference train in parallel like transformers and decode with a fixed-size state. Schmied et al. (2025) train xLSTM from scratch as the core of a large action modela single sequence model trained offline by next-action prediction on many control tasks for 432 simulated control tasks and compare it with a transformer of the same size. Hauzenberger et al. (2026) take the other route for language models: convert a pretrained transformer into an xLSTM hybrid by distillation, so the recurrent model inherits the transformer's training.

A recurrent backbone costs the same per step at any context length, so its advantage appears at long contexts and large batches, and in both papers quality held up on most of the benchmarks tested. Trained from scratch on control data at 16M to 206M parameters, xLSTM matched or beat the transformer, kept a flat step time of about 1.7 ms (read off a bar chart) and handled contexts at which the transformer ran out of memory. Distilled from 7B and 8B language models, with a 512-token sliding window of exact attention kept beside the recurrent memory, the student came within 5% of its teacher on at least half of the benchmarks it was scored on after about 6B to 21B training tokens, and decoded about twice as fast at 131K tokens. What it did not keep is exact retrieval from far back: on needle-in-a-haystacka test that hides a fact in a long stretch of irrelevant text and asks the model to retrieve it tests its accuracy was 0.15 at 4K tokens against the teacher's 1.0, so the contexts where the recurrent model is fastest are the ones where its quality was not shown.

Setting. The mLSTMthe matrix-memory cell of xLSTM: linear attention with an exponential input gate, a sigmoid forget gate and an output gate is linear attention with gates. For an input $x_t$ the vector entering the layer at token $t$ , with query, key and value projections $q_t,k_t,v_t$ , a feature mapa fixed or learned function $\phi$ applied to queries and keys before their inner product; the distilled student uses a softmax over the feature dimension $\phi$ , learned weights $w_i,w_f,W_o$ , the logistic sigmoid $\sigma$ , an input gate $i_t=\exp(w_i^\top x_t)$ , a forget gate $f_t=\sigma(w_f^\top x_t)\in(0,1)$ and an output gate $o_t=\sigma(W_ox_t)$ , it keeps a matrix memory and a normalizer of fixed size (stabilization of the exponential gate omitted): $\begin{aligned}S_t&=f_t\,S_{t-1}+i_t\,\phi(k_t)v_t^\top,\\z_t&=f_t\,z_{t-1}+i_t\,\phi(k_t),\\h_t&=o_t\odot\frac{S_t^\top\phi(q_t)}{z_t^\top\phi(q_t)}.\end{aligned}$ The forget gate lets the memory drop old content, and the exponential input gate lets one token write strongly. The memory parallelizes over the sequence for training and runs one token at a time for decoding. The sLSTMthe scalar-memory cell of xLSTM, an LSTM with exponential gating whose recurrent weights mix the hidden state across steps; it can track state but cannot be parallelized over time is the other xLSTM cell; xLSTM [7:1] mixes seven mLSTM blocks per sLSTM block and xLSTM [1:0] uses mLSTM only. The action model reads three tokens per timestep, the state, the return-to-gothe sum of rewards still to come in the episode; conditioning on a high target return asks the model for high-reward actions and the reward, through encoders shared across tasks (a CNNconvolutional neural network, here the image encoder used for Atari and Procgen for images, a linear layer otherwise), and predicts all action dimensions at once with one shared head over 2,066 classes, with each continuous dimension cut into 256 bins. Previous actions are left out of the input.

The distilled student. Every attention layer of the teacher becomes a hybrid with two branches that share the teacher's query, key and value projections: an mLSTM over the whole past, and sliding-window attentionsoftmax attention restricted to the most recent tokens, here 512, so its cache never exceeds the window (SWA) over the last 512 tokens plus four attention sinksthe first few tokens of the sequence, kept in the window permanently because attention heads put much of their weight on them. The mLSTM's output gate $o_t$ , here one sigmoid scalar per head computed from the concatenated query, key and value, mixes the two: $\begin{aligned}\hat h_t&=o_t\,\underbrace{\frac{S_t^\top\phi(q_t)}{z_t^\top\phi(q_t)}}_{\text{mLSTM, all past}}\\&\quad+(1-o_t)\,\underbrace{\mathrm{SWA}(q_t)}_{\text{last 512, 4 sinks}}.\end{aligned}$ Training runs in two stages, with an optional third. Stage I fits only the new parameters (feature maps and gates), layer by layer, to the teacher's attention outputs by mean squared error, on 655M tokens. Stage II trains everything end to end on 5B to 20B tokens, mixing next-token cross-entropy with a KL divergence $\mathrm{KL}(p\|q)=\sum_y p(y)\log\frac{p(y)}{q(y)}$ , here from the teacher's next-token distribution to the student's to the teacher, restricted to the teacher's 256 most likely tokens: $\begin{aligned}\min_\theta\sum_t\Big[&-\gamma\log p_\theta(y_t\mid x_{1:t})\\&+\beta\,\mathrm{KL}\big(p_T^{(256)}\,\big\|\,p_\theta^{(256)}\big)\Big],\end{aligned}$ with $\gamma=0.9$ and $\beta=0.1$ . The truncation means the teacher's targets can be computed once and stored, so the teacher is not needed during Stage II. Stage III trains one student per domain (math, STEM, code, chat, about 5B tokens each) from the same starting weights and takes a weighted average of their weights.

For a controller that needs a long history at a fixed rate, the action-model results say a recurrent backbone gives the same quality at a lower and flat per-step cost; at the short contexts used in training (50 timesteps) all backbones were well under 10 ms, so the gain is in memory and in what context lengths become possible. For language models, the distillation recipe converts an 8B transformer with a small fraction of a pretraining budget and keeps most benchmark quality, and the student's memory per sequence is fixed: in each layer and head, the mLSTM state plus 512 window tokens and 4 sinks. It should not yet replace a transformer where answers depend on exact recall of distant context, and the reported gains were measured in a research implementation; serving systems such as vLLM are built around a paged KV cache, and scheduling mixed layer types in them efficiently is open systems work, as the authors note.

Step 1 (action model)

At equal size, the recurrent backbones match or beat the transformer. The data are 3.4M trajectories (894M transitions) from 432 tasks in six simulated domains, Atari, Procgen, Meta-World, DMControl, Composuite and Mimicgen, collected from expert or RLreinforcement learning: training an agent online to raise its reward-trained agents; 37 further tasks are held out. xLSTM [7:1], xLSTM [1:0], Mamba and a GPT-2-style Decision Transformera transformer trained to predict the next action from the recent states, returns-to-go and rewards (Chen et al., 2021) were trained at 16M, 48M, 110M and 206M parameters for 200,000 updates with a context of 50 timesteps (150 tokens). Both recurrent backbones reach lower validation perplexitythe exponential of the mean cross-entropy of the predicted actions; lower is better than the transformer at every size, and higher normalized scoresreturns rescaled per task relative to the returns in the training data (relative to human players for Atari), then averaged over domains when the models are run in the environments; at 206M the score is about 0.78 for xLSTM against 0.73 for the transformer (approximate, read off the plot), and xLSTM pulls ahead of Mamba at the largest size. Per domain, xLSTM is best on three of six and tied on the rest. A pretrained 16M xLSTM fine-tuned on the held-out tasks beats one trained from scratch in most domains, so the backbone does not hurt transfer. The comparison gives the recurrent models twice the transformer's number of layers at the same parameter count; halving their layers, as in the timing tests, costs a little and still beats the transformer.

Step 2 (action model)

The recurrent backbone's step time does not grow with context, and the transformer's memory does. On an A100 at 206M parameters, timed while playing Atari Freeway (episodes of 8,192 steps, so 24,576 tokens, with the simulator included), xLSTM with the transformer's layer count takes about 1.7 ms per step (all timings in this step are approximate, read off the paper's bar charts) at batch size 1 at every context from 50 to 25,600 timesteps, against about 2.3 ms for the transformer with a KV cache and FlashAttentiona fused GPU kernel for exact attention that avoids writing the full attention matrix to memory; the transformer runs out of memory at 25,600 timesteps (76,800 tokens). At batch size 16 the transformer's step time more than doubles by 3,200 timesteps and it runs out of memory from 6,400, while xLSTM stays near 2 ms. At a context of 1,600 timesteps the transformer runs out of memory above batch size 64, where the recurrent models keep scaling. Mamba was slower than both at batch size 1 (about 4.8 ms), which the authors attribute to its kernels not supporting compilation. For xLSTM, throughput depends on the head dimension: at fixed model width, more and smaller heads raise it, because the memory updates cost about the width times the head dimension in total.

Step 3 (action model)

Longer context helps once previous actions are removed, and the sLSTM helps only where state must be tracked. With actions in the input, a transformer trained on all 432 tasks scores about 0.3 at every context length; without them its score rises from 0.3 with one timestep to 0.7 with 50 (one seed per length). In robotics data actions change smoothly, so a model that sees the last action learns to copy it, the copycat probleman imitation-learning failure in which a policy predicts the next action from the previous one and ignores the observation (Wen et al., 2020); its predictions become overconfident, and a poor first action is then repeated. Removing actions forces the model to use the observations, and only then does history help, which is the regime where a recurrent backbone's fixed step cost matters. The mLSTM-to-sLSTM ratio matters little on the 432 tasks, where mLSTM-only does best on average (16M models, one seed per ratio), but on Dark-Room, a partially observed grid world used to test in-context RLimproving at a new task across episodes kept in the context, with no weight updates (Laskin et al., 2022), xLSTM [7:1] scores highest on both training and held-out tasks; the authors attribute this to the sLSTM's ability to track statekeep an exact running summary that every step updates, such as the parity of a bit string or the position reached after a sequence of moves, which, by Merrill et al. (2024), transformers and Mamba cannot do in general.

Step 4 (distillation)

The student matches its teacher on half or more of the benchmarks within a small tolerance. Ratios of student to teacher score average poorly when some teacher scores are near zero and can hide large losses, so the paper counts benchmarks instead. With $\mathcal B$ a set of benchmarks and $A_S(b),A_T(b)$ the student's and teacher's scores, the win-and-tie ratethe fraction $C_\alpha$ of benchmarks on which the student scores at least $(1-\alpha)$ times the teacher at tolerance $\alpha$ and the summary $\alpha^*$ are $\begin{aligned}C_\alpha&=\frac{1}{|\mathcal B|}\sum_{b\in\mathcal B}\mathbf 1\big\{A_S(b)\\&\qquad\ \ge(1-\alpha)A_T(b)\big\},\\\alpha^*&=\inf\{\alpha:\ C_\alpha\ge1/2\},\end{aligned}$ and the paper calls distillation losslesswhy: the paper's definition, $C_0\ge1/2$ : the student matches or beats the teacher on at least half the benchmarks; it says nothing about how much is lost on the other half when $C_0\ge1/2$ . From the base Llama-3.1-8B, the student averages 74.4 on six multiple-choice understanding tasks against the teacher's 74.0 (LoLCATs, an earlier conversion to linear attention, 69.9), and 36.3 against 35.5 on seven generation tasks in math and code, where LoLCATs scores 1.9; $\alpha^*=0$ . From Olmo-3-7B, $\alpha^*=0.01$ over the 13 benchmarks pooled, but the generation average is 34.7 against 41.6, so the half of the benchmarks that $\alpha^*$ does not look at can carry real losses. The ablations, on 2.5B tokens, support each part: mLSTM beats plain linear attention, adding the window gives a large further drop in loss and the sinks a smaller one; a pure KL objective does worse than the 0.9/0.1 mix, full fine-tuning beats LoRAlow-rank adaptation: training only low-rank updates to the weight matrices, a cheaper alternative to updating all weights, and Stage I alone plateaus quickly.

Step 5 (distillation)

Training experts and averaging their weights beats one student trained on the mixture. From the instruction-tuned Llama-3.1-8B and Qwen-2.5-7B, four domain students of about 5B tokens each are merged with weights chosen on small validation sweeps; $\alpha^*$ is 0.02 and 0.05. A single student trained on the mixed data for the same 20B tokens does worse on the specialized benchmarks, for Llama MATH500 0.37 against 0.54, HumanEval 0.50 against 0.63 and IFEval 0.49 against 0.69; the authors attribute this to each expert spending all its updates on one domain. Merging helps instruction following most (IFEval for Llama) and hurts STEM questions most (GPQA), where the merged model falls below the STEM expert. Judged by GPT-5.1 on MT-Bencha set of multi-turn chat prompts whose answers a strong model grades from 1 to 10, both students score above their teachers (6.05 against 5.08 for Llama, 5.96 against 5.20 for Qwen), a result from one automatic judge. The earlier converted models compared, Mamba-in-Llama (which keeps half of its attention layers) and QRWKV7 (an RWKV-7 conversion of Qwen), recover much less on math and code.

Step 6 (distillation)

The student is two to four times faster at long contexts, where its recall fails. On one H100, both models implemented in the same library with compiled kernels, the student's prefillencoding the prompt before the first generated token; its time is the time to first token throughput is about twice the teacher's at batch size 1 and 65K tokens, with about half the time to first token; generating 131K tokens from no prompt takes about half the teacher's time and memory, with the student's memory flat; and with prompts and batch size 8 its generation throughput is up to about four times higher, while the teacher runs out of memory at larger batches. On needle-in-a-haystack retrieval the instruction-tuned teachers score 0.998 to 1.0 at 1K to 16K tokens, and the Llama student scores 0.69, 0.15, 0.08 and 0.02 at 1K, 4K, 8K and 16K (single needle; the Qwen student is similar). The student was distilled on 4K-token sequences with no long-context data, and fine-tuning the transformer itself without conversion also lost long-context ability, so the authors cannot say whether the fixed-size memory or the training data is the cause.

Terms

mLSTM: $S_t=f_tS_{t-1}+i_t\phi(k_t)v_t^\top$ , $h_t=o_t\odot S_t^\top\phi(q_t)/z_t^\top\phi(q_t)$ : gated linear attention, fixed-size state
sLSTM: scalar-memory xLSTM cell with recurrent mixing; tracks state but runs sequentially in time
xLSTM [7:1], [1:0]: seven mLSTM blocks per sLSTM block; mLSTM only
Large Recurrent Action Model: xLSTM backbone; 3 tokens per step (state, return-to-go, reward); shared head over 2,066 action classes; no actions in the input
Copycat problem: with previous actions in the input, the policy copies them; removing them made longer context help (0.3 to 0.7)
Latency (206M, A100): xLSTM about 1.7 ms per step at any context; transformer about 2.3 ms, out of memory at 25,600 timesteps (read off a bar chart)
Hybrid student layer: $\hat h_t=o_t\,\mathrm{mLSTM}+(1-o_t)\,\mathrm{SWA}$ over 512 recent tokens and 4 sinks
Distillation stages: I: layer-wise MSE to teacher outputs (655M tokens); II: 0.9 cross-entropy + 0.1 top-256 KL (5B to 20B); III: merge domain experts
Win-and-tie rate $C_\alpha$: share of benchmarks with $A_S\ge(1-\alpha)A_T$ ; $\alpha^*=\inf\{\alpha: C_\alpha\ge1/2\}$ ; lossless means $C_0\ge1/2$
Long-context recall: needle-in-a-haystack 1.0 for the teacher; 0.69, 0.15, 0.08, 0.02 for the student at 1K, 4K, 8K, 16K

Quiz

Why does the recurrent backbone's advantage over the transformer grow with context length and batch size?
Why does the distilled student keep a 512-token sliding window of softmax attention beside the mLSTM?
Four domain experts of about 5B tokens each were merged and compared with one student trained on the mixed data for 20B tokens. What did the comparison show?

Sources

Thomas Schmied, Thomas Adler, Vihang Patil, Maximilian Beck, Korbinian Pöppel, Johannes Brandstetter, Günter Klambauer, Razvan Pascanu, and Sepp Hochreiter. A Large Recurrent Action Model: xLSTM Enables Fast Inference for Robotics Tasks. ICML, 2025.
Lukas Hauzenberger, Niklas Schmidinger, Thomas Schmied, Anamaria-Roberta Hartl, David Stap, Pieter-Jan Hoedt, Maximilian Beck, Sebastian Böck, Günter Klambauer, and Sepp Hochreiter. Effective Distillation to Hybrid xLSTM Architectures. arXiv preprint, 2026.